[논문리뷰] U-Net: Convolutional Networks for Biomedical Image Segmentation

Paper

U-Net: Convolutional Networks for Biomedical Image Segmentation

---

Abstract

• 딥러닝 학습에는 수천 개 이상의 annotation 학습 데이터가 필요하다.
• context 포착을 위한 contracting path와 위치 추정을 가능하게 하는 expanding path로 구성되있다.
• 512X512 이미지를 분할하는데에 2015년 GPU로 1초가 채 걸리지 않는다.

---

Introduction

• 더 큰 모델과 더 많은 양의 데이터가 주를 이루었다.
• CNN의 주요 task는 단일 항목에 대한 분류이지만, 생물학 이미지에서는 분류 뿐 아니라 픽셀 단위의 분류가 필요하다.
  

  

• 이 네트워크는 pixel 단위의 class 예측이 가능하고, patch 단위의 학습으로 training data를 더 많이 학습할 수 있다.
• 그러나 이 방법은 두가지 단점을 가진다.
  1. 각각의 patch에 대해 개별적인 학습을 진행해야하며, patch들 간의 중복이 많기 때문에 학습에 시간이 많이 걸린다.
  2. localization과 context간의 trade-off 관계가 발생한다. (큰 patch는 많은 maxpooling layer가 필요하고, 작은 patch는 좁은 context만 확인한다.)
• 본 논문에서는 더 좋은 Fully Convolutional Network를 제안한다.
• 붙어있는 세포에 대한 경계를 포함하고 있기 때문에, 별도의 weighted loss 함수를 제공한다.

---

Network Architecture

• 모델의 좌측에 contracting path, 우측에 expansive path가 있다.
• contracting path는 전형적인 convolution network 구조를 따른다.
  • 반복적인 두번의 3x3 conv가 반복된다 (padding = 0)
  • ReLU와 2x2 maxpooling (stride=2)가 뒤따른다.
  • 각 step을 거치면서 channel의 수가 2배가 된다.
• expansive path는 upsampling으로 구성된다.
  • up-convolution이라 불리는 2x2 conv가 반복된다.
  • channel은 반으로 줄어들고, 3x3 conv가 반복된다.
  • cropping은 경계 픽셀에 대한 loss에 필요하다.
• 마지막 layer는 1x1 conv이다.
• 모델은 총 23개의 conv layer를 가진다.

---

Training

• SGD 방법으로 학습한다.
• unpadded이기 때문에, output image가 input image 보다 일정 너비만큼 작다.
• GPU를 최대한 사용하기 위해서, batch size는 줄이고 input size의 크기를 키웠다.
• momentum을 키워서, 이전 학습 방향을 따라가도록 했다.
• Energy function은 cross entropy loss와 마지막 feature map의 pixcel-wise soft-max를 결합하여 계산한다.
• Energy function
  

  

• weight map
  • w0 = 10, sigma = 5
  • 세포 사이의 거리가 가까우면, w(x)의 값은 커진다.
    

    

• unet은 가우시안 분포를 통해 가중치를 초기화 한다.
  • $\sqrt{2/N}$ (N = previous layer conv size * channels)

---

Data Augmentation

• Elastic deformation : 탄성변형, 외부의 힘이 제거되면 원래대로 돌아가려는 성질 (세포가 눌렸다가 다시 펴지는 모습을 상상)
  

  

---

Conclusion

• elastic deformation으로 적은 양의 데이터와 이상적인 학습시간을 가질 수 있었다. (NVidia Titan GPU 6GB 10hours)