[논문 리뷰] U-Net 꼼꼼한 리뷰

Abstract

Deep Neural Network의 성공적인 학습을 위해서는 많은 양의 annotated training sample들이 필요하다.

이 paper에서는 annotated sample들을 더욱더 효율적으로 활용하기 위하여 Data augmentation기반의 새로운 network 와 training strategy 에 대하여 소개한다.

해당 구조는 context(이미지의 내용)를 capture하기 위하여 수축하는 path와 정확한 localization을 위한 symmetric한 확장 path로 구성되어 있다.

이러한 network는 적은 양의 image들에 대하여 end-to-end로 학습이 가능하다.

Introduction

Convolutional Network는 보통 classification task, 즉, image의 output 결과가 single class label인 경우였다.

그러나, 많은 visual task에서 특히 biomedical image processing 에서 기대하는 output은 localization 즉, class label이 pixel로 표시되어야 했다.

이 paper에서는 “fully convolutional network”라는 이름의 elegant한(ㅋㅋ,,) 구조로부터 구상되었다. 우리는 이 구조를 수정하고 확장하여 적은 양의 이미지들에 대해서도 작동할 수 있도록 하였으며, 보다 더 정확한 segmentation을 예측할 수 있도록 하였다.

fully convolutional network 란, 기존의 classification network에서 출력층의 fc layer를 conv layer, 즉 upsampling layer로 대체한 것이다.

즉, fully convolutional network의 주요 아이디어는 네트워크를 점점 축소하여 class label만을 예측할 수 있었던 기존의 CNN 네트워크와 달리, 마지막 출력층에 upsampling layer (즉, pooling 연산자 대신 upsampling 연산자를 추가함)을 추가하여 객체의 위치정보를 픽셀 단위로 표시할 수 있도록 한것이다.

따라서, 이러한 layer는 output의 해상도가 증가하게 해준다.

localize를 위해서는 contracting path로부터 나온 고해상도의 feature들이 upsample된 output과 결합된다.

기존의 FCN 구조와 비교해서 u-net 구조의 수정된 부분은 upsampling 하는 부분에서 feature channel의 개수를 늘림으로써 context information을 higher resolution layer에 전파할 수 있도록 한것이다.

결과적으로, expansive path (확장하는 경로)는 contracting path(수축 경로)와 symmetric하게 구성되며, 이는 u-shaped 구조로 나타난다.

해당 네트워크는 fully connected layers(모든 뉴런이 연결되어있는 상태) 즉 2차원에서 1차원으로 평탄화를 수행하는 layer가 없다.

이 이유는, segmentation task의 output 인 segmentation map이 input image의 conext 정보와 위치 정보를 담은 pixel만을 포함하기 때문에 fc layer가 필요 없기 떄문이다.

Network Architecture

해당 구조는 왼쪽 (contracting path) 와 오른쪽(expansive path) 로 이루어져있다.

contracting path는 convolutional network의 전형적인 구조를 따라간다.

contracting path 구성

• 3x3 convolutions + ReLU 2번
• 2x2 maxpooling layer 1번

를 4번 반복

각 downsampling path에서 feature channel의 개수를 두배씩 늘린다. (그림에서 보면, 64 channel의 개수를 downsampling을 통해 128개의 feature channel로 늘린것을 볼 수 있다. )

그 밑에 쓰여져있는 숫자, 572x572, 570x570 등등은 feature map의 height, width를 의미하는데, unet구조는 valid padding, 즉 padding을 사용하지 않았기 떄문에 convolution을 거칠때마다 output의 크기가 줄어드는 것을 볼 수 있다.

단계	연산	크기변화
입력		572 x 572
conv1	3×3 valid conv	570 x 570
conv2	3×3 valid conv	568 x 568
maxpool	2x2 maxpool	284 x 284
conv3	3x3 valid	282 x 282
conv4	3x3 valid	280 x 280
…	…	…

여기서 잠시, feature channel의 개수가 무엇을 의미하는지 알아보자.

feature channel의 개수란, 다음 그림에서 처럼 convolution 연산에서 필터의 depth를 의미한다.

https://devkor.tistory.com/entry/합성곱-신경망

expansive path 구성

• contracting path에서 나온 feature map을 concatenation 시킴
• 2x2 up convolution 1번
• 3x3 convolutions + ReLU 2번

final layer 의 1x1 convolution은 semantic segmentation task 에 맞추어 공간 해상도를 유지하면서 채널의 수를 줄이기 위한 방법이다.

결과적으로 전체 네트워크는 23개의 convolutional layer를 가진다.

 

Skip Connection 

 

Unet 구조에서 또 하나 중요한 점은, skip connection이라는 구조이다. 이 구조 덕에 unet 모델은 세밀한 구조까지 모두 살릴 수 있는 모델이 되었다고 할 수 있다!

skip connection은 

Encoder(Contracting Path)에서 추출한 특징(feature map)을
Decoder(Expanding Path)에서 다시 활용하기 위해,
같은 레벨끼리 feature map을 직접 연결(concatenate) 해주는 구조라고 할 수 있다.

 

• Downsampling(Pooling)을 하면 해상도가 작아지면서 공간 정보(특히 경계 정보)가 손실된다.
• Decoder는 업샘플링으로 해상도를 복원하려고 하지만, 정확한 경계나 위치 정보는 이미 손실
• 따라서 Encoder의 고해상도 feature map을 Decoder에 건너뛰어(skip) 전달한다. 

즉, 저해상도의 의미 정보와 고해상도의 위치 정보를 모두 활용하고자 하는 모델이라고 볼 수 있겠다. 

 

Training

데이터 증강 : 해당 논문에서는 의료이미지를 훈련 데이터셋으로 사용하였기 때문에 데이터 수가 매우 부족하기 때문에 augmentation 기법이 필수적이다.

Augmentation

이 논문에서는 여러가지 augmentation 기법들 중 특히 elastic deformation을 사용하였다.

elastic deformation(탄성 변형) 이란?

이미지를 마치 고무처럼 비틀거나 흔드는 듯한 변형. 조직의 위치가 조금씩 다르더라도 모델이 잘 인식하도록 도와줌.

실제로는 gaussian filter를 사용하여 블러처리한다.

Loss Function

cross entropy 사용

$$ E = \Sigma_{x\in\Omega}w(x)log(p_{l(x)}(x)) $$

기호	의미
$$ Ω $$	이미지 상의 모든 픽셀 위치들의 집합 (즉, 전체 이미지 영역)
$$ x $$	개별 픽셀 위치
$$ l(x) $$	픽셀 x의 정답 클래스(label)
$$ p_{l(x)}(x) $$	픽셀 x에서 정답 클래스에 해당하는 확률값 (모델 출력 softmax 값)
$$ w(x) $$	픽셀 x에 부여된 가중치(weight) – 경계 강조 + 클래스 불균형 보정용

segmentation task는 절대적인 pixel 의 양으로 보았을 때, background는 매우 많고, target은 아주 적기 떄문에 weight에 가중치를 주었다.

$$ w(x) = w_c(x) + w_0 * exp( - \frac{(d1(x) + d2(x))^2}{2σ^2}) $$

두번째 항은 object 사이의 경계에 가까운 픽셀에 큰 값을 주고 있다. 즉, 가중치 기반 loss는 "배경은 너무 많으니까 덜 중요하게, 중요한 부분(장기, 경계 등)은 더 집중해서 배우자!"는 전략이다.