U-Net은 2015년 U-Net: Convolutional Networks for Biomedical Image Segmentation이라는 논문을 통해 발표되었다. Segmentation테스크를 수행하는 모델 중 하나로 성능을 많이 개선시켰다.


Abstract

Deep Networks를 학습하는데는 많은 학습 데이터가 필요하다는 사실은 모두가 알고 있다. 그래서 이 논문에서는 적은 샘플에서 Augmentation테크닉을 이용해 Network를 학습하는 전략을 소개했다. 적은 샘플에서도 ene-to-end학습이 가능하고 이전에 가장 좋은 모델이었던 sliding-window 방식보다 더 좋은 모델을 소개한다. 똑같은 미생물 사진으로 학습했을 때 큰 차이로 다른 모델의 성능을 뛰어넘었고, 더 빠르기까지 했다. 512x512 사이즈의 사진 1장을 처리하는데 GPU기준 1초보다 더 낮은 속도를 보였다.


Introduction

근래 2년간 CNN은 인식기술에 있어서 뛰어난 성능을 보였다. 물론 CNN이라는 개념은 오래전부터 있었지만 복잡한 모델을 학습하는데 필요한 데이터가 부족했고, 모델 자체의 크기를 감당하는데 어려움이 있었다. Krizhevsky et al.가 수많은 ImageNet 데이터셋 학습을 시도한 이후에 많은 발전이 있었다.


보통 CNN은 1가지 label이 있는 classification 테스크에서 사용이 되었는데, biomedical image를 다룰때는 이미지의 각 pixel이 라벨링되어있는 localization이 필요했다. 더 나아가 대부분의 학습 데이터는 biomedical과는 거리가 있었다. 실제로 Ciresan et al.가 sliding-window방식을 이용해 각 픽셀의 label를 예측하는 시도를 했었다. 이 network는 localization이 가능했고 학습 데이터보다 더 많은 양의 patch를 이용해 학습을 했다. 그 결과 ISBI2012에서 큰 격차로 우승을 했다.


하지만 Ciresan et al.의 방법에는 2가지 문제점이 있었다. 첫 번째는 한 이미지로부터 생성된 patch가 각각 분리되어 학습이 되었기 속도가 매우 느렸고 이미지 내에 중복되는 부분이 많았다. 두 번째는 localization의 정확도와 context(이미지 내의 정보를 인식하는 정도)의 trade-off관계가 있었다. 정보를 많이 인식하기 위해 patch의 사이즈를 키우면 그 만큼 max-pooling layer가 많이 필요하기 때문에 localization의 정확도가 낮아지고, 반대로 patch의 사이즈를 줄이면 context가 낮아진다.


이 논문에서는 “fully convolution network(FCN)”을 발전시킨 모델을 제시했다. 적은 수의 학습 데이터로도 작동하고 segmentation에서 더 좋은 정확도를 보이는 모델이다. FCN의 핵심은 pooling연산자가 upsampling 연산자로 대체함으로써 해상도를 높이는 것이었다. 또한 localization을 하기 위해서 각 단계별로 추출된 ‘해상도가 좋은’ features와 upsampled된 output을 결합했다. 이런 식의 구조로 인해 convolution layer들이 정보를 더 정확하게 학습할 수 있게 된다.



이 논문에서 제시한 모델의 중요한 부분은 FCN와 다르게 upsampling을 하는 부분에서도 channel이 많은 feature를 사용한다는 점이다. 그래서 각 대응하는 upsampling layer에 정보를 제공할 수 있다. 그래서 결과적으로 그림을 보면 U자 모양의 모델이 만들어진다. 또한 모델은 fc layer를 사용하지 않았다. 이 모델은 overlap-tile 전략을 사용했는데, 이미지의 구석부분에 CNN을 적용하기 위해 padding을 적용했다. Padding의 경우 Mirror방식을 사용했다. 이러한 방식은 큰 이미지에 적용할 때 매우 유용했지만, GPU의 한계로 특정 사이즈까지만 확대가 가능했다.


저저가 하려고 하는 task에서는 데이터가 매우 적었기 때문에 elastic deformation을 적용해 augmentation을 실시했다. 이 방법은 변형에 매우 invariance하도록 학습이 되었다. Biomedical segementation에서는 변형에 invariance한 것이 매우 중요한데, 왜냐하면 조직검사에 있어서 변형이 가장 common하기 때문이다.


세포를 segmentation할 때 challenge한 부분은 같은 class인 세포들이 붙어있을 때 그들을 서로 분리하는 것이다. 저자는 붙어있는 세포들의 경계를 더 잘 학습하기 위해 weighted loss를 사용했다. 대충 biomedical segmentation분야에서 우수한 성적을 보였고, 특히 ISBI cell tracking challenge 2015에서 우수한 성적을 보였음.


Network Architecture


모델은 2파트로 이루어져 있는데, 그림에서 왼쪽 부분인 contracting path(수축하는, downsampling)와 오른쪽 부분인 expansive path(확장하는, upsampling)이다. 특징은 다음과 같다.

  1. Downsampling
    • 단계별로 3x3 unpadded convolutions, ReLU 2개
    • stride가 2인 2x2 maxpooling layer 1개
    • Downsampling시 channel를 2배로 증가
  2. Upsampling
    • 단계별로 3x3 unpadded convolutions, ReLU 2개
    • 2x2 up-convolution
    • Upsampling시 channel 1/2로 감소
    • Contracting에서 생성된 feature map을 crop한 feature와 concatenation적용


마지막 단계에서는 1x1 convolution을 이용해 class개수의 map을 생성한다. 모델어 사용된 convolution layer는 총 23개이다. 만약 seamlee tiling을 하기 위해서는 이미지의 사이즈를 짝수로 사용해야한다.


Training

학습을 할 때는 stochastic gradient descent를 사용한다. Unpadded convolution 때문에 output 이미지의 사이즈가 input 사이즈보다 작다. 그리고 GPU메모리를 최대한 활용하기 위해 batch 크기를 키우는 것 보다는 input의 tile를 키우는 것을 지향한다. 또한 momentum값은 0.99로 사용하기 때문에 이전 샘플이 현재 샘플을 optimize하는데 큰 영향을 준다.


마지막 output부분에서는 pixel-wise soft-max를 사용해 예측을 진행하고, cross entropy loss function을 사용한다. Softmax의 경우 다음과 같이 정의를 했다.


\[p_{k}(\mathbf{x})=\exp \left(a_{k}(\mathbf{x})\right) /\left(\sum_{k^{\prime}=1}^{K} \exp \left(a_{k^{\prime}}(\mathbf{x})\right)\right)\]


$K$는 class의 개수를 나타내고 $\mathbf{x}$는 픽셀을 나타낸다. 그리고 $p_{k}(\mathbf{x})$는 maximum-function으로 가장 큰 값은 1이고 나머지는 0이다. 최종적으로 loss는 다음과 같다.


\[E=\sum_{\mathbf{x} \in \Omega} w(\mathbf{x}) \log \left(p_{\ell(\mathbf{x})}(\mathbf{x})\right)\]


$\ell$은 각 pixel의 true label을 나타내고, $w$는 세포들의 class별로 frequency가 다른 픽셀을 보상해주기 위해, 특히 border픽셀에 더 큰 가중치를 주기 위한 가중치값이다. 이때 $w$는 사전에 ground truth segmentation dataset을 이용해 미리 계산을 한 후에 적용했다. $w$의 계산식은 다음과 같다.


\[w(\mathbf{x})=w_{c}(\mathbf{x})+w_{0} \cdot \exp \left(-\frac{\left(d_{1}(\mathbf{x})+d_{2}(\mathbf{x})\right)^{2}}{2 \sigma^{2}}\right)\]


$w_c$는 class frequency를 맞춰주기 위한 weight map이고, $d_1$은 x픽셀에서 가장 가까운 세포의 border까지의 거리를 나타내고, $d_2$는 x픽셀에서 2번째로 가까운 세포의 border까지의 거리를 나타낸다.


깊은 신경망에서는 다양한 cnn layer와 path가 존재하기 때문에 초기값은 매우 중요하다. 만약 초기값을 제대로 정해주지 않으면 특정 노드만 과도하게 큰 activation값이 출력될 우려가 있다. 따라서 각 feature map은 분산을 1로 설정을 해야한다. 논문에서 제시한 모델은 가우시안 분포를 사용하고 표준편차를 $\sqrt{2/N}$로 설정함으로서 실현할 수 있다. 여기서 N은 한 개의 뉴런에서 오는 node의 개수를 말한다. 예를 들어 3x3 Filter와 64개의 feature map이 있다면 $N$의 값은 $9\cdot64=576$이 된다.


Data Augmentation


만약 학습 데이터가 부족할 때 모델이 invariance하고 robustness하게 학습되기 위해서는 augmentation이 필수이다. 변형이 자주 발생하고 그레이 스케일인 미생물 이미지에서는 소위 robustness라고 말하는 shift와 rotation invariance에 강해야한다. 이 경우에는 random elastic deformation augmentation이 가장 좋은 방법이라고 한다.


Experiments

논문에서는 U-Net을 3가지 다른 segmenation task에 적용을 했다. 첫 번째 task는 electron microscopic recordings이다. 위에서 Overlap-tile을 소개할 떄 사용한 이미지가 그 예시이다. 또한 보조자료로 해당 task에 대한 모든 결과를 제시했다. 데이터는 EM segmentation challenge인 ISBI2012에서 제공했다. 512x512사이즈를 가진 30개의 이미지가 있었고 각 이미지에서 세포는 흰색, 세포 주의는 검은색으로 라벨링되어잇었다. 물론 테스트 데이터의 segmentation map은 공개되지 않았다. 예측한 결과를 주최측에 전달하여 평가를 받았다. 평가 지표로는 각각 ‘warping error’, ‘rand error’, ‘pixel error’가 있었다. U-Net은 이러한 결과에 대해 매우 좋은 성능을 보였다.



이 결과에서 중요한 점은 이전에 모델이었던 sliding-window방식보다 월등히 좋다는 것이다.



U-Net은 light microscopic 이미지에도 좋은 성능을 보였다. 이 데이터는 ISBI cell tracking challenge 2014와 2015이다. 먼저 ‘PhC-U373’데이터셋은 Glioblastoma-astrocytoma U373 세포를 포함하고 있고, 35개의 학습 데이터로 구성되어 있다. 이 데이터로는 IOU를 측정했는데, 다른 모델들보다 월등히 좋았다. 결과는 위 사진에 나와있다.


Conclusion

결론적으로 U-Net이 매우 다른 biomedical segmentation applications에 매우 좋은 성능을 보였다. 또한 Elastic deformation 덕분에 적은 양의 데이터로도 좋은 성능을 보일 수 있었다.


읽어주셔서 감사합니다.(댓글과 수정사항은 언제나 환영입니다!)