※ 해당 글을 작성하게 된 이유,
U-Net 은 Semantic Segmentation task 에서 기본이 되는 모델입니다.
참고로, Diffusion model ( DDPM ) 에서 U-Net 구조를 채택합니다.
(논문 그대로 해석한 것이 아닌 제 나름의 방식대로 해석해서 작성했음을 알리며 참고 바랍니다.)
Olaf Ronneberger, Philipp Fischer, Thomas Brox, U-Net: Convolutional Networks for Biomedical Image Segmentation, MICCAI(2015)
Summary
- FCN을 사용해 Semantic Segmentation task 를 수행
- Data augmentation 을 사용해 적은 데이터를 가지고도 좋은 성능을 보일 수 있음을 증명
1. Introduction
해당 논문에서 ISBI 2012년도 challenge 에서 우수한 성능을 보였던 네트워크의 2가지 결점을 말합니다.
- 하나, 각 patch 별로 계산이 수행되기에 속도가 느리면서 overlapping patch 가 있어 redundancy 가 큽니다.
- 둘, localization accuracy 와 context 간 성능에서의 trade-off 가 존재합니다.
▶ 큰 patch 는 max-pooling 을 많이 수행함으로써 네트워크가 좀 더 다양한 context 를 바라보게 하는 데 도움을 주지만 localization accuracy 를 줄이는 요소가 됩니다. 반면, 작은 patch 는 네트워크로 하여금 작은 영역의 context 만을 보게 합니다.
▶ 따라서 이 둘은 서로 trade-off 관계를 가지게 되는 것입니다.
( 위 내용은 FCN 에서도 비슷하게 언급되었으니 참고 바랍니다. )
2. Network Architecture
저자는 FCN 의 구조를 참고 및 수정함으로써 적은 training images 로 정확한 segmentation 결과를 산출할 수 있도록 했습니다.
- 많은 수의 feature channel 을 그대로 사용했습니다.
- 이는 더 큰 resolution 에 context information 을 전파하는 역할을 합니다.
- 결과적으로 contracting path 와 expansive path 는 대칭 구조를 지니게 됩니다.
- contracting path 는 전형적인 convolutional network 의 구조를 가지며 image 의 context 를 지니도록 합니다. 또한, expansive path 는 upsampling 한 feature map 을 context 와 결합함으로써 정확한 localization 을 수행하도록 합니다.
- 네트워크는 어떠한 fully connected layer 를 가지지 않기에 input image 에 전체 context 주입이 가능합니다.
▶ Strategy 1. Overlap-tile strategy
border 영역에서의 pixel 값을 예측하기 위한 목적으로 input image 를 mirroring 함으로써 크기가 큰 이미지를 네트워크에 적용하도록 하는 기법입니다.
- 이는 이미지의 해상도가 GPU 메모리 용량에 의해 제약받지 않도록 도움을 줍니다.
▶ Strategy 2. Data augmentation
- 적은 양의 training data 로 충분한 학습을 할 수 있도록 합니다.
- 네트워크는 augmentation 을 통한 변형을 통해 불변성을 학습합니다.
- 변형을 가하는 것은 실제 tissue 에서의 흔한 변형이기에 simulation 을 하는 효과를 불러일으킵니다.
3. Training
- unpadded convolution 연산으로 인해 output 이 input 보다 일정 크기 작게 나옵니다.
- overhead 최소화 & GPU 사용 최대화하기 위해 batch 를 1로 (= 작게) 두고 input image 를 크게 두었습니다.
- 이에 따라, high momentum 값을 이용해 이전에 학습되었던 많은 양의 training sample 들이 현재의 optimization 에 관여하도록 했습니다.
- Energy function 은 최종적인 feature map, cross entropy loss function, pixel-wise soft-max 를 결합해 계산됩니다.
$$ E=\sum_{x\in\Omega}^{}w\left(x\right)log\left(p_{l\left(x\right)}\left(x\right)\right) $$
- Soft-max $$ p_{k}\left(x\right)=exp\left(a_{k}\left(x\right)\right)/\left(\sum_{\acute{k}=1}^{K}exp{\left(a_{\acute{k}}\left(x\right)\right)}\right) $$
- Weight map 은 border 를 잘 구분해 cell segmentation 하도록 학습에 관여합니다. ( Figure 3 참고 ) $$ w\left ( x \right ) = w_{c}\left ( x \right ) + w_{0} \bullet exp\left ( -\frac{\left ( d_{1}\left ( x \right )+ d_{2}\left ( x \right )\right )^{2}}{2{\sigma }^{2}} \right ) $$
