[논문 리뷰] U-Net: Convolutional Networks for Biomedical Image Segmentation
U-Net:
Convolutional Networks for Biomedical Image Segmentation
Abstract
딥러닝 네트워크의 경우에 Sample에 대한 의존도가 크다. 이 연구에서도 Data Augmentation 학습 방법을 소개하고, U-Net이라는 네트워크를 소개한다. U-Net은 2개의 Part로 나뉠 수 있는데, Context를 파악하는 Contracting Path와 Localization의 정확도를 올려주는 Expanding Path이다. 이 두 가지 Path는 서로 대칭적으로 존재한다. 이 모델의 경우 아주 적은 Sample로도 학습이 잘 되었고, ISBI Challenge에서 이전 우승 모델보다 높은 성능을 보였다.
1. Introduction
본 논문에서는 CNN의 성공이 Training Set의 양이 커지면서 생긴 제한적인 이유라고 말한다. 이전까지는 CNN은 Classification을 위해 많이 사용되었다. 하지만 생물학 분야의 영상 처리에서는 Localization이 중요하였고, Semantic Segmentation의 중요도가 높았다. 하지만 생물학에 대한 Sample의 개수가 1000개밖에 되지 않는 것이 다수였다. 이전에도 이러한 연구가 있었지만, 두 가지 결점이 존재했다. Pixel 별 정확도가 중요하기 때문에 매우 느렸다. 그리고, Localization 정확도와 Context 정보 간의 Trade-off가 존재하였다. 이를 해결하기 위해 본 논문에서는 Fully Convolutional network를 소개한다.
위 사진은 모델의 구조도이다. 보이는 것처럼 Context 정보를 얻기 위한 Contractinig Path와 Localization을 위해 Resolution을 키우는 Expanding Path가 존재한다. 이 모델의 중요한 변화는 Upsampling Part에 Feature들이 큰 Channel을 가지고 있다는 점이다. 이는 Context 정보가 Higher Resolution으로 잘 전달될 수 있도록 한 것이다. Fully Connected Layer를 사용하지 않고 Convolution Layer만 사용했다는 특징이 있다. 또한, Border 영역의 Pixels을 잘 예측할 수 있도록 Missing Context를 Input Image를 Mirroring 하여 Extrapolate 하는 방법을 사용하였다. 큰 Image를 처리하기 위해서 Tiling Strategy를 사용하였다.
사용 가능한 Sample이 적기 때문에 Data Augmentation 방법을 사용하였다. 이때는 Elastic Deformation을 사용하여 Invariance를 학습하였다. 이러한 방법은 생물학 분야에서 많이 사용되는데, 대부분의 조직들은 탄성 변형을 하기 때문이다. 또한, Weighted Loss를 적용하여 Border 쪽을 더 잘 학습하도록 하였다.
2. Network Architecture
위에서 소개했듯이 모델의 구조는 두 파트로 나뉜다. 왼쪽은 Contracting Path, 오른쪽은 Expansive Path이다. Contracting Path는 전형적인 CNN 구조이다. 3x3 convolution으로 이루어져 있고, ReLU Activation Function을 사용한다. Downsampling을 할 때는 2x2 Max pooling을 사용한다. Downsampling이 될 때는 Channel의 크기가 두배가 된다. Expansive Path는 2x2 convolution을 통해 Upsampling을 한다. Upsamling 시에는 Channel이 절반이 된다. Contracting Path에서 나오는 Crop 된 Feature Map은 Expansive Path의 Feature Map과 Concatenation 되며, 2개의 3x3 Convolution과 ReLU를 거치게 된다. 이 Crop 된 Feature Map은 Border Pixel을 더 잘 학습하기 위해 사용된다. 마지막 Layer의 1x1 Convolution은 64개의 Feature Channel에서 Class개수로 Mapping 하기 위해 사용된다. 총 23개의 Convolution Layer가 사용되었다.
3. Training
SGD Optimizer를 사용하였고, unpadded Convolution을 이용하였다. GPU Memory 사용을 극대화하고 Overhead를 줄이기 위해, 큰 입력의 경우 Tiling을 하여 mini-batch로 쌓아서 사용하였다. Momentum은 0.99로 큰 값을 사용하였는데, 이는 이전 Training Sample에서 본 것들이 현재 Optimization Step에 영향을 많이 주도록 하기 위함이다. Softmax를 통해 마지막 Feature map에서 확률을 계산하여 Cross Entropy Loss Function과 함께 사용하였다.
위의 식은 Cross Entropy에 Weight map을 곱한 것이다. w는 Weight map으로 중요한 Pixel을 더 강조하기 위한 것이다. 각 Ground Truth마다 Class의 Frequency를 고려하여 Weight Map을 미리 계산하여 사용하였다. Separation Border는 Morphological Operation이라는 컴퓨터 비전 알고리즘을 통해 계산할 수 있다. Weight Map은 다음과 같은 식으로 계산하였다.
Wc는 Class Frequency의 Balance를 맞추기 위해서 사용한 것이고 d1은 근처 cell의 border로의 거리이고, d2는 두 번째로 가까운 Cell의 border 거리이다. w0와 sigma는 상수로 각각 10과 5로 실험하였다.
딥러닝에서 모델의 Layer가 커지면서 좋은 Weight Initialization 방법이 중요해졌다. 여기서는 Gaussian Distribution으로 Weight를 Initialize 해주었다.
1. Data Augmentation
Data Augmentation은 모델의 성능과 Robustness를 위해서는 필수적인 요소이다. 특히 Sample이 적을 때는 더 중요하다. 여기서는 Shift와 Rotation을 사용하였고, 특히 Random Elastic Deformation을 사용한 것이 Key Concept이다. Dropout Layer를 contracting Path에 추가함으로써 Data Augmentation과 같은 기능을 하도록 하였다.
4. Experiments
U-Net은 3가지 Segmentation Dataset에 사용하였다고 말한다. 첫 번째 Task는 여기서 진행한 미생물의 neuronal structure segmentation Task이다. EM segmentation challenge에서 Warping Error에 대해 가장 좋은 성능을 보였다.
두 번째는 Cell Segmentation Task이다. ISBI Cell Tracking Challenge 2014와 2015에서 PhC-U373 Dataset에 대해 가장 높은 성능을 보였고, 세 번째 Dataset인 DIC-HeLa Dataset에 대해서도 2등과 확연한 차이를 보이며 우승하였다.
5. Conclusion
U-Net은 생물학 Segmentation에 좋은 성능을 보였고, Elastic Deformation을 이용한 Data Augmentation의 영향을 크게 받았다. 이것은 추후 실험에서 적은 양의 사진으로 학습하는데 크게 기여하였다.
U-Net은 앞으로도 많이 보이게 될 구조이고, 특히 Segmentation 모델을 관찰하다 보면 U-Net과 비슷한 구조를 한 모델이 많이 보일 것이다.
질문은 언제나 환영입니다:)