Paper Details
Title: U-Net: Convolutional Networks for Biomedical Image Segmentation
Authors: Olaf Ronneberger, Philipp Fischer, Thomas Brox
Conference: Medical Image Computing and Computer-Assisted Intervention (MICCAI) 2015
Year of Publication: 2015
Link: https://arxiv.org/abs/1505.04597
Key Focus:
This paper presents U-Net, a convolutional neural network designed for biomedical image segmentation. The architecture is characterized by a symmetric U-shaped structure, consisting of a contracting path to capture context and an expansive path to enable precise localization. U-Net relies heavily on data augmentation, including elastic deformations, to maximize the utility of limited annotated datasets. It employs weighted loss functions to address challenges such as separating touching objects of the same class. The network achieves state-of-the-art performance in various biomedical tasks, including ISBI 2012 EM segmentation and ISBI 2015 cell tracking challenges, significantly outperforming previous methods. The implementation is efficient, processing a 512x512 image in under a second on a modern GPU.
💡Key point: U-net = 생물 의학 이미지 분할을 위한 딥러닝 아키텍처
📄 Summary of the Paper
1. Introduction
Deep convolution network는 여러 시각적 인식 작업에서 효과를 보였지만, 사용 가능한 학습 데이터셋의 크기와 네트워크 크기의 한계를 가지고 있었다. Convolution networks는 일반적으로 분류 작업에 사용되는 네트워크로 입력 이미지에 대해 단일 클래스 레이블을 출력하는 역할을 한다. 그러나 생물 의학 이미지 처리를 포함한 많은 이미지 처리 작업에서 각 픽셀마다 클래스 레이블을 할당하는 localization이 요구되었다. 동시에, 생물 의학 이미지 처리 작업의 경우에는 수천 개의 학습 이미지를 학습하는 것이 매우 어려웠다. 이에 따라 Ciresan et at.의 'Deep neural net- works segment neuronal membranes in electron microscopy images' 연구에서는 슬라이딩 윈도우 설정에서 픽셀 주변의 로컬 영역(패치)을 입력으로 사용해 각 픽셀의 클래스 레이블을 예측하는 네트워크를 학습하는 방식을 선보였다. 해당 네트워크는 localization이 가능하며, 패치의 양이 학습 이미지보다 훨씬 많아 데이터 양의 제약을 완화한다는 장점이 있었고, ISBI 2012 EM 분할 챌린지에서 큰 격차로 우승하는 성과를 거두었다.
하지만, Ciresan et al. 이 소개한 네트워크는 다음의 두 가지 단점을 가졌다.
- 네트워크를 각 패치마다 별도로 실행해야 하므로 속도가 느리다.
- localization의 정확도와 맥락 사용 간에 트레이드오프가 존재한다.
(큰 패치는 더 많은 맥스 풀링 계층을 필요로 하기 때문에 localization 정확도를 낮추고, 작은 패치는 적은 맥락만 볼 수 있다.)
본 논문에서는 완전 컨볼루션 네트워크(fully convolutional network)에 기반을 둔 아키텍처를 제안하며, 이를 수정 및 확장해 소량의 학습 이미지에서도 작동하며 더 정밀한 분할을 제공하도록 설계하였다.
핵심적인 아이디어는 일반적인 수축 네트워크를 업샘플링 연산자로 대체된 연속적인 계층으로 보완하는 것으로, 이러한 계층은 출력의 해상도를 증가시킨다. 또한, localization을 위해 수축 경로의 고해상도 특징을 업샘플링된 출력과 결합하였다. 연속적인 컨볼루션 계층은 이 정보를 기반으로 더 정밀한 출력을 학습 가능하다.
2. Network Architecture
앞서 언급했듯이 U-Net의 네트워크 구조는 수축 경로와 확장 경로로 구성된 대칭적 U자형 아키텍처를 기반으로 한다.
(네트워크는 총 23개의 convolution layer를 가진다.)
- 수축 경로 (전형적인 convolution networks의 구조를 따른다.)
- 3x3 컨볼루션(패딩 없음)과 ReLU를 두 번 반복 적용
- 2x2 맥스 풀링(stride 2)으로 다운샘플링
- 다운샘플링마다 피처 채널 수를 두 배로 증가
- 확장 경로
- 업샘플링 후 2x2 업컨볼루션으로 피처 채널 수를 절반으로 감소
- 수축 경로에서 잘린 피처 맵과 병합
- 3x3 컨볼루션과 ReLU를 두 번 반복 적용
- 출력 처리
- 1x1 컨볼루션으로 최종 분류
이와 같은 대칭적 설계는 맥락 정보와 세부 정보 간의 균형을 유지하는 역할을 하며 동시에 높은 해상도와 정밀도를 제공한다.
3. Training
U-Net은 입력 이미지와 해당 분할 맵을 사용하여 Caffe 기반 확률적 경사 하강법(SGD)으로 훈련된다. 패딩 없는 컨불로션 사용으로 인해 출력 이미지가 입력 이미지보다 작아지는 문제가 발생하는데, 이를 보완하기 위해 입력 타일 크기를 최적화한다. GPU 메모리를 최대한 활용하고 오버헤드를 최소화하기 위해, 큰 입력 타일을 선호하고 배치 크기를 단일 이미지로 줄인다. 이에 따라 높은 모멘텀(0.99)을 사용하여 이전에 학습한 많은 샘플이 현재 최적화 단계의 업데이트에 영향을 미치도록 하였다.
에너지 함수는 최종 특징 맵에서 픽셀 단위 소프트맥스를 적용하고 크로스 엔트로피 손실 함수와 결합하여 계산된다.
소프트맥스(Softmax)는 다음과 같이 정의된다.
- $a_k(x)$: 위치 $x$에서 특징 채널 $k$의 활성화 값
- $K$: 클래스 수
- $p_k(x)$: $k$가 가장 높은 활성화를 가질 때 1에 가까워지고, 다른 값은 0에 가까워짐.
크로스 엔트로피는 각 위치에서 실제 클래스 레이블과의 편차를 다음과 같이 페널티화한다.
- 훈련 중 특정 픽셀에 더 높은 중요도를 부여하기 위해 설계된 가중치 맵
훈련 데이터에서 특정 클래스 픽셀의 빈도 차이를 보완하고 접촉하는 세포들 사이의 작은 분리 경계를 학습하도록 강제하기 위해, 각 정답 분할에 대해 가중치 맵을 미리 계산한다. 분리 경계는 형태학적 연산을 사용하여 계산되며, 가중치 맵은 다음과 같이 계산된다.
- 클래스 빈도 균형을 위한 가중치 맵
- $d_1(x)$: 가장 가까운 세포 경계까지의 거리
- $d_2(x)$: 두 번째로 가까운 세포 경계까지의 거리
- 실험에서는 $w_0 = 10$과 $\sigma \approx 5$ 픽셀로 설정함
딥러닝 모델에서 가중치 초기화는 매우 중요한 요소이다. 초기화가 적절하지 않을 경우, 일부 네트워크 부분에서 과도한 활성화가 발생하고, 다른 부분은 거의 기여하지 못할 수 있기 때문이다. 이상적으로는 네트워크의 각 특징 맵이 대략적으로 단위 분산을 가지도록 초기화해야 한다. 본 연구의 아키텍처에서는 가중치를 평균 0, 표준편차 $\sqrt{2/N}$ 를 가지는 가우시안 분포에서 샘플링하여 초기화하였다.
3.1. Data Augmentation
데이터 증강은 적은 수의 학습 샘플로 네트워크가 불변성과 견고성을 학습하는 데 필수적인 기법으로, 많은 수의 학습 데이터셋을 얻기 어려운 생물 의학 이미지 처리 분야에서 빈번하게 사용된다. 또한, 현미경 이미지의 경우 이동 및 회전 불변성과 변형 및 회색 값 변화에 대한 견고성이 필요하다.
주요 증강 기법으로는 랜덤 탄성 변형, 회전, 이동 등이 있는데, 본 연구에는 랜덤 탄성 변형(random elastic deformations)을 핵심으로 다루고 있다. 랜덤 탄성 변형은 특히 조직 변형을 모델링하는 데 유용하며, 3x3 그리드에서 가우시안 분포를 따르는 변위 벡터를 생성하여 구현된다. 또한, 드롭아웃 레이어를 활용하여 네트워크의 일반화 능력을 향상시키는 추가적인 증강 효과를 제공한다.
4. Experiments
본 연구에서는 U-Net을 다음의 이미지 처리 작업에 적용한 결과를 제시하고 있다.
1. EM 분할 챌린지 (ISBI 2012)
- 데이터: 초파리 배신경줄기 단면 (512x512 크기, 전자현미경 이미지 30개)
- 결과:
- 워핑 에러: 0.000353 (새로운 최고 기록 달성)
- 랜드 에러: 0.0382 (기존 슬라이딩 윈도우 기반 네트워크보다 우수)
- 픽셀 에러: 0.0611
- 의의: 후처리 없이도 최고 성능을 달성하며, 데이터 특화 알고리즘을 능가함.
2. 세포 추적 챌린지 (ISBI 2014/2015)
(1) PhC-U373 데이터셋
- 데이터: 위상차 현미경 이미지, 35개 주석 데이터 포함
- 결과:
- 평균 IOU 92%
- 두 번째 최고 알고리즘(83%) 대비 우수
(2) DIC-HeLa 데이터셋
- 데이터: 차동간섭현미경(DIC) 이미지, 20개 주석 데이터 포함.
- 결과:
- 평균 IOU 77.5%.
- 두 번째 최고 알고리즘(46%)보다 월등히 높은 성능 달성
5. Conclusion
U-Net 아키텍처는 다양한 생물의학 분할 응용에서 매우 우수한 성능을 발휘함을 확인할 수 있었다. 탄성 변형을 활용한 데이터 증강 덕분에 소량의 주석 이미지만으로도 학습이 가능하며, 소요되는 학습 시간 역시 기존 네트워크에 비해 크게 단축되었다.
