U-Net: Convolutional Networks for Biomedical Image Segmentation

📄U-Net: Convolutional Networks for Biomedical Image Segmentation

written by Olaf Ronneberger, Philipp Fischer, and Thomas Brox

---

Introduction

지난 시간 동안 Deep Convolutional Networks(DCN)는 많은 시각 인식 작업에서 큰 성과를 거둠. DCN의 성공은 주로 대규모의 데이터셋과 네트워크 크기의 한계로 인해 제한되었으나 Krizhevsky 등이 8개 층과 수백만 개의 parameters를 가진 대형 네트워크를 ImageNet 데이터셋에 대해 감독 학습하여 성과를 냄

전형적인 DCN은 분류 작업에 사용되어 이미지에 a single class label을 output으로 도출함. 그러나 생물의학 이미지 처리에서는 각 픽셀에 label을 할당해야 하는 경우가 많음. Ciresan 등은 슬라이딩 윈도우 방식을 사용해 픽셀 주변의 국소 영역을 입력으로 받아 각 픽셀의 class label을 예측하는 네트워크를 훈련시킴. 그러나 이 방식은 느리고 중복이 많으며, 큰 patch에서는 문맥을 많이 사용할 수 있지만 위치 정확도가 떨어지는 문제가 있음

본 논문에서는 fully convolutional network를 기반으로 한 a more elegant architecture를 제안함. 해당 architecture는 훈련 이미지가 적어도 작동하며, 더 정확한 segmentation을 제공함. 주요 아이디어는 수축 경로에서 고해상도 feature를 upsampled된 output과 결합하여 정확한 output을 학습하는 것. 또한 data augmentation을 통해 훈련 데이터를 늘리고 변형에 대한 불변성을 학습하게 함

더하여 접촉하는 객체를 분리하기 위해 가중 손실 함수를 사용함. 제안된 network는 다양한 생의학 segmentation 문제에 적용 가능하며, EM stack에서 신경 구조의 세분화 및 ISBI 2015의 세포 추적 챌린지에서 우수한 성과를 보임

---

Network Architecture

Network architecture는 contracting path(왼)와 expansive path(오)로 이루어짐

contracting path는 전형적인 convolution network architecture를 따름. 두 개의 3x3 convolution(no pading)과 각 conloution 후에 정류된 ReLU, 그리고 2x2 max pooling으로 구성됨. 다운샘플링 단계마다 feature channel의 수가 두 배로 증가

expansive path의 각 단계는 feature map을 업샘플링하고 2x2 convolution(up-convolution)을 통해 feature channel 수를 절반으로 줄인 후, contracting path의 대응되는 부분에서 잘라낸 feature map과 결합함. 이후 두 개의 3x3 convolution과 각 convolution 후에 ReLU가 적용됨. 잘라내기는 각 convolution에서 경계 픽셀이 손실되기 때문에 필요함. final layer에서는 1x1 convolution을 사용하여 각 64-component feature vector를 원하는 class 수로 mapping함. network는 총 23개의 convolution layers를 가짐

output segmentation map의 원활한 tiling을 위해 모든 2x2 max pooling 연산이 x 및 y 크기가 짝수인 layer에 저굥ㅇ되도록 input tile size를 선택하는 것이 중요함

---

Training

Network는 Caffe의 Stochastic Gradient Descent(SGD) 구현을 사용하여 input image와 해당 segmentation map으로 훈련됨. padding이 없는 convolution으로 인해 output image는 input보다 경계 너비만큼 작아짐. GPU 메모리를 최대한 활용하기 위해 큰 batch 크기보다는 큰 input tile을 사용하고, patch는 단일 이미지로 줄임. 높은 momentum(0.99)을 사용하여 이전 훈련 샘플들이 현재 최적화 단계의 업데이트를 결정하게 함

energy 함수는 최종 feature map에 대한 pixel 단위 sofmax와 cross entropy loss 함수로 계산됨. softmax는 각 pixel 위치에서 k번째 feature map의 activation을 기반으로 정의됨. cross entropy는 각 위치에서 예측된 class 확률과 실제 label의 차이를 벌칙으로 사용함

각 ground truth segmentation에 대해 weight map을 사전 계산하여 특정 클래스의 픽셀 빈도를 보상하고, 서로 접촉하는 cell 사이의 작은 분리 경계를 학습하도록 함. 분리 경계는 형태학적 연산ㅇ르 사용하여 계산됨. weight map은 class 빈도를 균형 있게 하기 위한 wc와 가장 가까운 cell 경계와 두 번째로 가까운 cell 경계까지의 거리를 사용하는 식으로 계산됨

deep netwrok에서는 weight initialization이 매우 중요. 네트워크의 일부가 과도한 활성화를 제공하거나 기여하지 않게 되는 것을 방지해야 함. 네트워크의 각 feature map이 대략 단위 분산을 가지도록 초기 가중치 설정. 이를 위해 가중치는 표준 편차가 sqrt(2/N)인 정규 분포에서 추출되며, 여기서 N은 한 뉴런의 들어오는 노드 수를 나타냄. 예를 들어, 3x3 컨볼루션과 이전 층의 64개 특징 채널에서는 N = 9 * 64 = 576

Data Augmentation

적은 수의 훈련 샘플로 네트워크에 필요한 invariance과 robustness를 가르치는 데 필수적

현미경 이미지를 대상으로 할 때, 주로 이동 및 회전 invariance와 변형 및 회색 값 변동에 대한 robustness가 필요함. 특히, random elastic deformations of the training samples은 매우 적은 수의 주석이 달린 이미지로 segmentation network를 훈련하는 데 중요한 개념임. 3x3 격자에서 랜덤 변위 벡터를 사용하여 매끄러운 변형을 생성함. 변위는 표준 편차가 10픽셀인 가우시안 분포에서 샘플링되며, 픽셀 단위 변위는 3차 스플라인 보간법을 사용하여 계산됨. contracting path 끝에 drop-out layer를 사용하여 추가적인 암묵적 data augmentation을 수행함

---

Experiments

본 논문에서는 U-Net을 세 가지 다른 segmentation 작업에 적용함

1. 전자현미경 기록에서 신경 구조의 세분화
   • 데이터 세트는 ISBI 2012에서 시작된 EM 세분화 챌린지에서 제공
     • 30개의 이미지(512x512 픽셀)와 해당 세포(흰색) 및 막(검은색)의 주석이 달린 세분화 맵으로 구성
   • U-Net은 전처리나 후처리 없이 0.0003529의 워핑 에러와 0.0382의 랜드 에러를 달성하여 최고 기록을 경신
     
2. 학 현미경 이미지에서의 세포 세분화
   • ISBI 세포 추적 챌린지 2014 및 2015의 일환
   • 첫 번째 데이터 세트 : PhC-U373
     • 위상차 현미경으로 기록된 글리오블라스토마-아스트로사이토마 U373 세포들
     • U-Net은 평균 IOU가 92%로 두 번째로 좋은 알고리즘(83%)보다 우수한 성과를 보임
   • 두 번째 데이터 세트 : DIC-HeLa
     • 차등 간섭 대비 현미경으로 기록된 HeLa 세포로 구성
     • U-Net은 평균 IOU가 77.5%로 두 번째로 좋은 알고리즘(46%)보다 뛰어난 성과를 보임
       

---

Conclusion

U-Net architecture는 다양한 생의학 segmentation 응용 분야에서 매우 좋은 성능을 발휘함. elastic deformations을 이용한 data augmentation을 통해 소수의 주석이 달린 이미지로도 훈련이 가능하며, NVidia Titan GPU(6GB)에서 10시간 정도의 훈련 시간만 필요로 함. U-Net architecture는 더 많은 작업에 쉽게 적용될 수 있을 것임.