U-Net

분야

Network Architecture

리뷰 날짜

2020/08/17

본 포스트는 제가 휴먼스케이프 기술 블로그에 먼저 작성하고 옮긴 포스트입니다.

본 포스트에서는 biomedical segmentation 분야에서 인용 수 16000+을 육박하는 논문에 대해서 리뷰하려고 합니다. 논문의 제목은 다음과 같습니다.

“U-Net: Convolutional Networks For Biomedical Image Segmentation”

논문에 대한 내용을 직접 보시고 싶으신 분은 이곳을 참고하시면 좋습니다.

Network Architecture of U-Net

Objective

논문에서 목적으로 하는 바는 biomedical image segmentation입니다. 자세히 설명하자면, 전자 현미경 등으로 관찰한 이미지에서 특정 세포(ex. 신경세포, 암세포)를 배경으로부터 분류하는 작업입니다. 본래 segmentation이라 하면 object detection을 포함하는 개념으로 다양한 종류의 객체들을 분류해내고 찾아내는 것이지만, 논문에서 초점을 맞추고자 하는 바는 이미지 속에서 다양한 세포를 분류해 내는 것이 아닌 단순히 한 종류의 객체들을 배경과 완전히 분리하는 것입니다.

구조적 관점의 목적은 U-Net architecture 의 마지막 layer의 output으로 산출되는 segmenation map을 학습을 통해 얻어낸다고 보시면 됩니다.

Comparison with Prior Study

논문에서 주로 비교의 목적으로 꺼내는 논문이 Ciresan et al.(이하 [1])입니다. [1]의 architecture는 EM segmentation challenge at ISBI 2012에서 우승을 차지했을만큼 성공적이었지만, 크게 두 가지의 문제점이 존재했습니다.

첫 번째는 localization(이미지의 세부 특성 반영)을 위해서 image를 patch단위로 분리하여 window-slide 형태로 학습시켜서 데이터의 중복이 매우 많으며, 학습 시간이 길었다는 점입니다.

Window-slide

Window-slide는 위 그림처럼 image를 patch로 나누어 학습시킬 때 공통된 부분이 굉장히 많고, 이에 따라 patch 개수도 많아집니다. 효과적인 localization을 위해서는 좋은 방법일 수 있으나, 학습 속도 측면에서 좋은 방법이 아닙니다.

두 번째는 localization과 context간의 trade-off가 존재하여 두 특성을 모두 챙기기 어려웠다는 점입니다.

Patch size가 크면 feature data가 context(이미지의 개략적인 패턴 반영)에 대한 내용은 포함하기 용이하지만 localization에 대한 내용을 포함하기는 어렵습니다. 반면에, patch size가 작으면 localization에 대한 내용은 포함하기 용이하지만 context에 대한 내용은 포함하기 어렵습니다.

Methods

앞서 제시한 선행연구의 문제점을 해결하기 위해서, 논문에서는 크게 두 가지의 개선방안과, 목적을 효과적으로 달성하기 위한 설계를 진행합니다.

첫 번째 문제점을 해결하기 위해서 논문에서는 Overlap-Tile Strategy를 사용합니다.

Overlap-Tile Strategy

Overlap-Tile Strategy는 위 그림처럼 파란색 범위의 input image patch를 넣었을 때 노란색 범위의 segmented image가 나온게 되는 상황에서, output segmented image가 겹치는 범위 없이 나오도록 input image patch의 설정 범위를 조절하는 방법입니다. 우측의 두 노란색 범위가 등장하는 output segmented image일 때 설정해야 할 두 input image patch의 범위는 우측의 두 파란색 범위와 같습니다. Overlap이라는 단어는 padding이 존재하지 않는 논문의 architecture 특성상 input에 비해 output의 size가 줄어들기 때문에 나타나는 input 설정의 겹치는 특성을 일컫는 말입니다.

이와 같이 patch size를 설정하여 과도하게 많은 patch 개수를 줄이는 전략을 사용했습니다. 여기에 더불어 최종적으로 빈 부분이 되어지는 가장자리 부분은 mirroring extrapolation(가장자리를 거울 대칭)을 이용해서 채우게 됩니다.

U-Net Architecture두 번째 문제점을 해결하기 위해서 논문에서는 U-Net Architecture를 구현했습니다. U-Net Architecture가 기존과 가장 다른 특이한 점은 Contracting Path, Exansive Path, 그리고 Contracting Path로부터 Expansive Path로 주어지는 Skip Connection이 존재한다는 점입니다.

Detail U-Net Architecture

Contracting Path는 일반적인 deep learning architecture에서 보이는 바와 같이 image로부터 feature를 extracting하는 부분입니다. Expansive Path는 segmentation map(각 pixel에 대한 binary classification map)을 다시 만들어내기 위해서 다시 upsampling 하는 과정입니다.

여기까지는 deep learning의 기본적인 설계인 feature extraction과 목적성에 맞게 upsampling하는 과정으로 볼 수 있지만, U-Net에서는 특별하게 작은 patch size에서도 feature가 context에 대한 정보를 포함할 수 있게 하기 위해서 Skip Connection이라는 것을 추가합니다.

Contracting Path에 존재하는 각 layer단계에 존재하는 일부 feature를 잘라서 Expansive Path의 feature에 concatenate해주어 결과적으로 upsampling 단계에서 늘어난 channel 수가 context에 대한 정보를 가질 수 있게 설계를 진행한 것입니다.

이렇게 전반적으로 U 형태를 그리는 architecture를 구현하여 U-Net이라는 이름을 가지게 됩니다.

Touching Object Classification논문에서 직면했던 문제 중 하나는 겹쳐져 있는 객체들의 효과적인 분류였습니다. 이를 해결하기 위해 논문에서는 학습을 위한 Cross Entropy를 정의할 때 가장자리와 가까운 pixel에 더욱 가중치를 두어 학습을 진행시켰습니다. 이에 대한 자세한 설명은 뒤에서 이어서 진행하도록 하겠습니다.

Touching Object Classification

Training

논문에서 weight 학습을 위한 기준이 되는 energy function은 pixel-wise soft-max function과 cross-entropy loss function으로 계산합니다.

먼저 pixel-wise soft max function은 다음과 같습니다.

p_k(X) = \frac{e^{a_k(X)}}{\sum_{k'=1}^K e^{a_k(X)}}

K는 전체 label의 수, k는 그 중 특정 label을 지정하는 변수이며, 결과적으로 마지막 output channel의 index라고 볼 수 있습니다. a_k(X)는 X의 위치(pixel x,y)에 있는 activation의 값입니다.

이 soft max function의 설계에서 기대하는 것은 특정 pixel의 위치 X에 대해 올바른 k로 학습되었을 때 p_k(X)가 1에 가까운 값에 수렴하고, 다른 경우엔 0에 가까운 값으로 수렴하는 것입니다.

다음으로 cross-entropy loss function은 다음과 같습니다.

E = \sum_{X\in\Omega} \omega(X)\log(p_{l(X)}(X))

w(X)항은 가중치에 대한 항목이며 l(X)는 위치 X에 대한 ground truth label입니다. 전체 function의 형태는 흔히 알고 있는 cross-entropy의 형태와 매우 유사하며 다른 점이 있다면 가중치에 대한 항목입니다. 각 위치에 대한 ground truth channel의 soft max function의 logarithm을 합산한 형태로 볼 수 있으며, soft max function 결과가 제대로 예측될 수록 에너지가 커지는 형태임을 확인할 수 있습니다.

더불어 앞서 생략한 가중치에 대한 설명을 간단히 드리자면, 다음과 같은 수식으로 정의했습니다.

\omega(X) = \omega_c(X)+\omega_0e^{-\frac{(d_1(X)+d_2(X))^2}{2\sigma^2}}

w_c(X)는 각 training data 별로 class frequency(labeling된 pixel의 비율들)이 다른 것을 조절해 주는 weight map입니다. d_1(X)는 가장 가까운 cell과 border까지의 거리이고 d_2(X)는 두 번째로 가까운 cell과 border까지의 거리입니다. 여기서 주목할 점은 border까지의 거리가 가까운 위치에 있는 pixel들이 기여하는 에너지에 더 가중치가 높다는 점입니다. 논문에서는 이 방법을 통해서 인접한 객체들의 경계를 확실히 구별하기 위한 과제였던 touching object classification 문제를 해결하려고 했습니다.

마지막으로, 논문에서는 초점을 맞추었던 문제의 특성상 학습에 많은 데이터를 구할 수 없었기 때문에 data augmentation을 진행했습니다. 그 방법으로 elastic deformation을 사용했는데 각 pixel을 gaussian distribution을 가지는 random displacement vector로의 변환을 이용해서 데이터 수를 인위적으로 늘렸다는 정도만 언급하고 넘어가도록 하겠습니다.

Experiments & Results

논문에서 제시한 실험과 그 결과는 논문 연구진들이 출전했던 EM segmentation challenge와 ISBI cell tracking challenge의 결과였습니다.

Ranking on the EM segmentation challenge 2015, sorted by Warping Error

Warping error의 관점에서 논문에서 제시한 U-Net이 가장 좋은 결과를 보여주었습니다. 이는 topological 관점에서의 equality를 따진 error이며 rand error보다 경계의 shifting에 대해서 민감한 error입니다. 마지막으로 Pixel Error는 가장 일반적으로 생각할 수 있는, pixel-wise로 제대로 segmentation이 이루어졌는지를 판단하는 error입니다. 자세한 error에 대한 설명은 이곳을 참고하시면 좋을 것 같습니다.

더불어 수치로 또 제공한 것이 IOU(Intersection Over Union)입니다. 이는 ground truth와 결과를 비교함에 있어 교차영역을 합영역으로 나눈 값입니다. 즉, 얼마나 실제와 동일하게 겹쳐져 있느냐를 측정하는 척도로 보시면 됩니다.

Segmentation results(IOU) on the ISBI cell tracking challenge 2015

위 결과처럼 각 cell별로 본 논문이 구현한 architecture가 가장 높은 IOU 수치를 보여주었습니다. 아래 그림은 각각의 cell을 실제로 tracking한 모습입니다.

Results on ISBI cell tracking challenge 2015

Conclusion

이것으로 논문 “U-Net: Convolutional Networks For Biomedical Image Segmentation”의 내용을 간단하게 요약해보았습니다. 실제로 이 논문은 논문의 결론 부분에서 U-Net 이 다양한 biomedical segmentation application에 적용될 수 있을 것이라고 언급한 것에 기대라도 하는 듯한 높은 인용치를 선보였습니다.

개인적으로 딥러닝이 biomedical 분야에서 현재 진행 중인 기여보다 앞으로 더 많은 것들을 할 수 있을 것 같음을 느낄 수 있었던 논문이 아니었나 싶습니다.