U-Net: Convolutional Networks for Biomedical Image Segmentation

본 내용은 논문 https://arxiv.org/abs/1505.04597 을 참고하여 작성하였습니다.

해당 논문은 2015년 5월 18일에 게재되었으며 논문의 저자는 Olaf Ronneberger, PhilippFischer, Thomas Brox입니다. 세 분 모두 독일의 프라이부르쿠 대학(University of Freiburg) 소속입니다.

논문의 핵심 내용은 주제에서 비춰지듯이 Biomedical Image segmentation을 위한 convolutional network에 대한 내용입니다. 

그럼 이제 부터 논문에 대해서 살펴보겠습니다.

저자는 지난 몇 년간 deep convolutaional network는 visual recognition에서 큰 업적을 이뤄왔습니다. 그렇지만 training에 있어서 training data set size, network size 등의 문제로 더 높은 성과를 올리는데 제약 사항이 있었음을 언급하였습니다.

그리고 이러한 문제를 supervised learning을 통해 해소할 수 있었음을 언급하였습니다. 즉, U-Net도 supervised learning을 사용하는 network임을 알 수 있습니다. 

Convolutional networks의 전통적인 사용은 하나의 영상에 대한 하나의 class만 label하는 classification 이었습니다. 그렇지만 생물의학 분야에서는 하나의 이미지에 다수의 class를 label해야만 했습니다. 

그래서 Ciresan과 그의 동료들은 픽셀의 local region(patch)을 주어 window sliding 방식으로 학습을 시키는 network를 구축했습니다. 그리고 이 network은 ISBI 2012, EM segmentation challenge에서 우승하였습니다.

저자는 물론 이 network가 좋은 기록을 내었지만 두 가지 결함이 존재함을 지적하였습니다.

첫 째는 각각의 patch에 대해 network가 개별적으로 돌기 때문에 속도 측면에서 다소 느리며, patch가 서로 overlapping 되었을 시 낭비가 심하다 입니다. (어떤 측면에서 낭비가 있는지는 언급하지 않았지만 여러개의 patch가 overlapping되었을 시 정확한 prediction을 내지 못함을 지적하는 듯 합니다.) 

둘 째는 localization accuracy와 use of context 사이의 trade-off가 존재함입니다. 하나의 이미지에 대해서 좋은 segmentation 효과를 보기 위해선 localization과 context는 중요한 자원입니다. 그렇지만 이 둘이 비례할 수 없다는 것을 지적하였습니다.

U-Net은 fully convolutional network(FCN)를 기반으로 하여 구축되었습니다. 저자는 적은 데이터를 가지고도 더욱 정확한 segmentation을 내기 위하여 FCN의 구조를 수정하였습니다. 다음은 U-Net의 구조입니다.

모양이 U자 형태를 띄는 것을 확인할 수 있습니다.

network의 중심을 기준으로 왼쪽을 Contracting path, 오른쪽을 Expansive path라고 정의했습니다.

그림에서 보이는 blue box는 multi-channel feature map을 의미하고 그 옆에 작은 숫자들은 x,y 사이즈를, 그 위에 숫자는 channel의 개수를 의미합니다. 그리고 각각의 화살표는 색 마다 서로 다른 operation을 의미하는데 그 의미는 그림 우측 하단에 명시되어있습니다.

자 이제 그럼 그림을 보면서 architecture의 내부 특징과 그 과정에서 사용된 Overlap-tile기법에 대해 살펴보겠습니다.

 

U-Net의 architecture에서 Contracting path는 이미지의 context를 포착할 수 있도록 돕고 Expansive path는 feature map을 upsampling하고 이를 contracting path에서 포착한 feature map의 context와 결합하여 더욱 정확한 localization을 하는 역할을 합니다. 그리고 이것이 U-Net의 main idea입니다.

기존 FCN에서와 다른 중요한 점이 있늕데 upsampling과정에서 feature channel의 수가 많다는 것입니다. 이는 network가 context를 higher resolution successive layer에게  전파할 수 있음을 의미합니다. 이 netwrok는 fully connected layer를 포함하지 않으며 오직 각각의 convolution의 valid part만을 사용합니다. valid part는 full context가 들어있는 segmentation map입니다. 이러한 기법은 overlap-tile 기법을 사용하여 seamless segmentation을 가능하도록 합니다.

그림에서 보이듯이 blue tile과 yellow tile이 서로 overlap되어 있음을 확인 할 수 있으며, blue area에 기반하여 yellow area의 segmentation을 prediction 했음을 확인 할 수 있습니다. 그 이후로 missing data는 mirroring을 통하여 extrapolation 합니다.

 

이러한 overlap-tile strategy를 적용함으로써 얻을 수 있는 혜택은 image resolution이 GPU memory에 대해 자유로워진다는 것입니다.

또한 U-Net은 적은 데이터로 충분한 학습을 하기 위해서 data augmentation을 사용합니다. 

이제 전체적인 논문의 U-Net 이론적 설명은 이정도에서 끝날 것 같습니다. 

지금부터는 논문의 구현 mechanism과 실험, 마지막으로 결론으로 이 글을 마무리하도록 하겠습니다.

좀 더 쉬운 이해를 위해 위에서 보았던 그림을 다시 가져오겠습니다.

Contracting path는 일반적인 convolutional network를 따르며 downsampling을 위한 stride2, 2x2 max pooling연산과 ReLU를 포함한 두 번의 반복된 3x3 unpadded convolutions 연산을 거칩니다. (3x3conv → ReLU → 2x2 max pooling → 3x3conv → ReLU → 2x2 max pooling) 

그리고 downsampling과정에서 feature map channel의 수는 두배로 증가시킵니다.

그리고 이를 Expansive path에서 2x2 convolutions를 통해 upsampled된 feature map과 contracting path에서 cropped 된 feature map과의 concatenation후에 ReLU 연산을 포함한 두 번의 3x3 convolutions 연산을 거칩니다.

(2x2 conv for upsampling → concatenation with cropped feature map in contracting step → 3x3 conv → ReLU → 3x3 conv → ReLU)

cropping이 필요한 이유는 각각의 convolution 과정에서 있어서 border pixel의 손실 때문입니다. 

마지막 layer는 1x1 convolution 연산을 하는데 이는 64개의 component feature vector를 desired number of classes에 mapping 하기 위함입니다.

다음은 그렇다면 이 model을 가지고 어떠한 방식으로 training을 진행하는지 알아보겠습니다.

input image와 그게 상응하는 segmentation map은 SGD와 함께 network를 학습하기 위해 사용됩니다. convolution 연산이 padding을 거치지 않기 때문에 input image ㅡ기 보다 output image의 크기가 더 작아집니다. 

overhead와 GPU memory의 사용을 극대화하기 위하여 큰 batch size보다는 큰 input tiles를 선호합니다. 따라서 저자는 높은 momentum(0.99)를 사용해서 과거에 진행했던 training sample이 현재의 update에 더욱 관여하도록 하였습니다.

Energy function은 최종적인 feature map과 corss entropy loss function, pixel wise soft-max를 결합하여 계산됩니다.

soft-max의 식은 다음과 같습니다.

Cross entropy는 각각의 position에서 penalize를 부과합니다. 식은 다음과 같습니다.

논문의 저자는 각각의 Ground truth segmentation의 weight map을 서로 다른 frequency를 대처하고 network에게 작은 separation borders를 학습시키기위해 사전에 weight map을 계산합니다.

(위 그림에서 (a)는 현미경으로 찍은 HeLa cells의 raw image, (b)는 raw image에 segmentation ground truth를 overlay, (c)는 생성된 segmentation mask, (d)는 network에게 border pixels를 학습시키기 위해 pixel-wise loss weight을 가진 map입니다.)

앞에서 말씀드렸던 separation border는 morphological operation을 통해 계산되어지며, weight map이 계산되어지는 식은 다음과 같습니다.

논문의 저자는 자신들의 실험에서 w0를 10으로 표준편차(σ)를 약 5pixel로 설정했다고 합니다.

저자는 deep network에서 초기 weight에 대한 good initialization의 중요성에 대해 언급합니다. 

이상적으로는 initial weights들은 각각의 feature map이 대략적인 unit variance를 가지도록 설정되어야 한다고 합니다.

저자는 U-Net은 

다음과 같은 표준편차를 가지는 Gaussian distribution으로 부터 initial weights를 도출해낼 수 있다고 합니다. N을 계산하는데에 있어 추가적인 설명을 붙였는데 만약 이전 layer에서 3x3 convolution 연산과 feature channel의 개수가 64개라면 N은 3x3x64 = 576으로 계산될 수 있다고 합니다.

기나긴 training에 대한 설명을 마쳤고 이제 실험과 결론을 말씀드리겠습니다.

저자들은 자신들이 구축한 network의 performance를 증명하기 위해 총 세 가지 실험을 거쳤습니다. 첫 번째 실험은 the segmentation of neuronal structures in electron microscopic recordings입니다. 

training data set은 각각의 이미지는 annotation된 cell은 흰색, membrane은 검은색으로 표시된 ground truth segmentation map과 함께 30개의 512x512 이미지를 사용하였다고 합니다. 

평가 결과는 predicted membrane probability map을 organizer에게 보냄으로써 이루어졌는데, 평가는 probability map을 10가지 different level 로 분류하고 "warping error", "rand error", 그리고 "pixel error"을 계산함으로써 이뤄졌습니다.

위 그림에서 3위를 차지한 IDSIA는 위에서 언급했던 Ciresan과 그의 동료들이 sliding window 기법으로 만든 segmentation network입니다. 그들의 기록보다 저자들이 좋은 성적을 거뒀음을 확인할 수 있습니다.

두 번째와 세 번째 실험은 PhC-U373, DIC-HeLa data set을 사용하였습니다.

(a)는 PhC-U373 dataset의 raw image이며, (b)는 segmentation result(cyan mask)와 ground truth(yellow border)입니다, (c)는 DIC-HeLa dataset의 raw image이고 (d)도 마찬가지로 mask와 gt가 결합된 결과입니다.

다음은 실험 결과입니다.

위 실험에서 사용된 평가지표는 average IOU(intersection over union)입니다. 표를 통해서 1위(u-net)과 2위의 성능 차이가 확연하게 드러남을 확인할 수 있습니다.

마지막으로 결론을 언급하고 글을 마치겠습니다.

결론적으로 저자는 자신들이 구축한 U-Net으로 몇 가지 실험을 거쳤고 이를 토대로 U-Net의 성능이 아주 좋음을 언급합니다. 그리고 그 공을 data augmentation으로 돌렸습니다. data augmentation 덕분에 매우 적은 train data와 적은 시간(Nvidia Titan GPU(6GB)를 사용하여 10시간 training)을 소비하고도 다른 networks보다 월등히 좋은 성능을 내었기 때문입니다.

이상으로 U-Net 논문 발표를 마치겠습니다.

감사합니다.