U-Net: Convolutional networks for Biomedical Image Segmentation

U-Net: Convolutional networks for Biomedical Image Segmentation

U-Net: Convolutional networks for Biomedical Image Segmentation

Introduction

이미지 세그멘테이션에선 픽셀 단위로 정답을 맞춰야 하기 때문에, 한 픽셀을 중심으로 주변 영역(Patch)을 잘라내서 CNN 모델에 입력으로 넣는 방식을 이전에 사용했다. 각 영역을 모델의 입력으로 넣기 때문에 하나의 이미지만으로도 여러 영역으로 나눌 수 있어 특히나 부족했던 의학 데이터 수가 마치 뻥튀기될 수 있었다.

하지만, 이미지 한 장에 대해서가 아니라, 한 이미지로부터 파생된 여러 영역에 개별로 모델을 돌리기 때문에 그만큼 계산량과 시간이 많이 요구되었다. 거기다 각 영역끼리 겹치는 부분으로 인해 중복된 연산 또한 적지 않았다. 한마디로 비효율적이란 것이다. 그리고 예측의 정밀도(Localization)와 맥락(Context)의 Trade-off가 있었는데, 입력이 되는 영역을 크게 잡으면 픽셀 주위의 정보가 많아지지만, 해상도가 뭉개지면서 정확한 픽셀 위치를 잡기 어렵고, 반대로 영역을 크게 잡으면 주변 맥락이 적어져, 이게 세포막인지 배경인지 구별하기 어려워진다.

이에 UNet은 FCN(Fully Convolutional Network)의 철학으로 여러 영역으로 나눠서 학습하는게 아니라 이미지 한장 전체를 학습할 수 있게끔하여, 앞선 두 문제를 해결하고자 하였다.

Architecture

모델의 구조는 위와 같다. Downsampling을 하는 Contracting path, Upsampling을 하는 Expanding path가 서로 대칭을 이루며 U자의 구조를 가지고 있고, Upsampling에서도 기존 FCN과 다르게 특징 채널 수를 크게 잡아 맥락 정보가 높은 해상도 레이어로도 잘 전파되게끔 했다.

또한 Downsampling에서 이미지 크기가 줄어들며 생기는 픽셀들의 위치 정보 손실을, 이미지 크기를 줄이기 전에 정보가 담긴 feature map을 Skip connection으로 upsampling 과정에서의 feature map에 채널 방향으로 합쳐서(Concatenation) 학습함으로써 극복하고자 함을 알 수 있다. 이때 모든 컨볼루션 레이어는 padding을 사용하지 않아 연산마다 이미지 크기가 줄어드는 것을 볼 수 있는데, 때문에 skip connection시 Contracting path의 출력의 크기를 줄인 후에 합친다.

padding을 사용하지 않는 이유는, 학습에 사용되는 이미지의 가장자리 부분은 중앙에 비해 주변 맥락이 부족하므로, 이미지 크기를 줄여나가며 주변 맥락이 완벽히 존재하는 중앙 부분만 학습을 진행하기 때문이다.

모델 구조에서 업샘플링시를 제외한 각 컨볼루션 레이어 이후에는 ReLU 활성화 함수가 따르고, 다운샘플링 시에는 이미지 크기를 절반으로 줄이는 대신 특징 채널을 두배로 키우며, 컨볼루션으로 인한 업샘플링에선 이를 반대로 수행한다. 최종적으로 $1 \times 1$ 컨볼루션 레이어어를 통해 특징 채널의 수를 예측할 클래스 수 만큼의 특징 채널로 매핑한다.

Implementation

Unpadded 컨볼루션으로 인해 이미지의 크기가 계속해서 줄어드는데, 원하는 출력 크기를 얻기 위해서는 입력 이미지의 크기가 그만큼 커야 한다. 당시에 GPU 메모리의 한계로 논문에선 배치 크기를 1로 줄여서, 즉, 이미지 한장씩 학습하면서, 최대한 이미지의 넓은 맥락 정보를 학습하는 것이 낫다고 판단하였다.

손실 함수는 모델 출력에 픽셀 별 Softmax에 가중치맵과 CrossEntropy를 통해 계산되었다. 즉, 다음과 같다.

\[L = -\sum_{x \in \Omega} w(x)log(p_{\ell}(x)(x))\]

이때 $x \in \Omega$는 이미지의 각 픽셀을 나타내고, $p_{\ell}(x)(x)$는 픽셀 $x$에서 True 클래스의 Softmax 출력 값이다. 다음과 같다.

\[p_{k}(x) = \frac{\exp(a_k(x))}{\sum_{k'=1}^{k}\exp(a_{k'}(x))}\]

이때 $a_k(x)$는 픽셀 $x$에서 채널(클래스) $k$의 활성화 값이다.

여기서 가중치 맵 $w(x)$와 softmax값을 CrossEntropy 해주는데, 이때 가중치맵은 학습 이전에 미리 계산되어 이미지의 각 픽셀이 세포 경계선에 얼마나 가깝게 붙어있는가를 나타낸다.

\[w(x) = w_{c}(x) + w_{0} \cdot \exp \left(-\frac{(d_{1}(x) + d_{2}(x))^{2}}{2\sigma^{2}}\right)\]

이때 $w_{c}$ 클래스간 불균형을 해결하기 위한 가중치 맵으로, 예를 들어 이미지에서 배경과 세포의 비율이 $9:1$로 크게 치우처진 경우에 모델이 모든 픽셀이 배경이라고 예측하도록 하는 것을 방지하기 위해 세포 픽셀 위치에 가중치를 더 높게 두는 것이다. $d_1$, $d_2$는 각각 각 픽셀에서 첫번째, 두번째로 가까운 세포의 테투리까지의 거리를 나타낸다.

그러면 $w(x)$는 두 세포 사이에 낀 경계선에 포함된 픽셀의 경우, $d_1$, $d_2$가 모두 작으므로, $\exp$ 항이 최대값으로 수렴하게 되고, 다른 픽셀보다 $w_0$만큼 가중치가 커지게 된다. 세포 내부 한가운데 있는 픽셀의 경우, $d_{1}$은 0에 가까워지지만 $d_2$는 커진다. 따라서 $d_1 + d_2$이 커져 $\exp$ 항의 값이 작아지며 해당 픽셀의 가중치가 비교적 적어진다.

즉, 두 세포 사이 경계선 픽셀의 가중치를 크게 설정되도록 하고, 이를 손실 함수에 CrossEntropy에 넣음으로써 모델이 이 세포끼리의 경계선만큼은 확실하게 학습하게끔 한다.

위 사진은 UNet 모델 학습 후 세포 세그멘테이션 결과인데, 빨간 영역이 모델이 세포라 예측한 영역, 파란 영역이 실제 세포 영역이며, 보라색 영역은 빨간 영역과 파란 영역이 겹치는 부분, 즉, 모델이 알맞게 예측한 영역이다. 사진을 보면 어느정도 오차는 있지만 적어도 세포간 맞닿은 경계 부분 만큼은 가중치 맵으로 인해 어느정도 명확히 구분하는 것을 볼 수 있다.

UNet은 padding이 없는 네트워크 구조와 단일 이미지 배치의 특성상 초기 가중치 설정이 중요하므로, He Initialization을 사용하였다. 이는 각 레이어의 가중치들이, 이전 레이어의 입력 노드 수 $N$을 기반으로 표준 편차를 $\sqrt{\frac{2}{N}}$로 설정하여 가우시안 정규분포에서 추출된 값을 사용하는 것이다.

UNet은 부족한 데이터셋에 있어서 데이터 증강을 사용하였다. 위치 이동, 회전, 명암의 변화뿐만 아니라, 탄성 변형(Elastic Deformation)을 적극적으로 활용하였는데, 이는 $3 \times 3$ 그리드 위에서 가우시안 분포로 무작위로 변위를 준뒤, Bicubic Interpolation을 통해 픽셀들을 부드럽게 왜곡한다. 이를 통해 실제 세포의 물리적 변형 형태 (말랑말랑한 느낌?)을 모사하여 적은 데이터셋으로도 모델이 세포 고유의 형태의 불변성을 학습할 수 있게끔 하였다.

논문에서 Contracting path의 끝에 Dropout 레이어가 또한 데이터 증강의 효과를 낸다고 했는데, 이는 내 생각으론 Dropout 레이어가 학습할 때 무작위 뉴런의 신호를 무시하여 내부 연산 경로가 바뀌게끔 하는데, 마치 미세하게 다른 모델을 번갈아가며 학습하듯이 앙상블의 효과를 내는 것이다. 모델 구조를 무작위로 흔들어 데이터가 늘어난 것과 같은 효과를 내는 이유라고 생각한다.

찌라시

UNet은 Downsampling과 Upsampling의 대칭을 이루는 U자 구조와 Skip connection을 통해 기존 픽셀의 정밀도와 픽셀 주변 맥락 정보의 trade-off를 극복하였고, 가중치 맵을 모델이 세포가 맞닿은 경계선까지 학습할 수 있게끔 하였다.

또한 알맞은 데이터 증강을 통해 기존 부족한 데이터셋의 환경에서도 모델이 형태적 불변성을 학습하게끔 하였다.