[논문 리뷰] U-Net ( Convolutional Networks for Biomedical Image Segmentation )

Introduction

  논문은, 기존의 classification tasks 들과는 달리, 다른 visual tasks들 특히, biomedical image processing의 경우 각 pixel에 class label 예측이 필요하다고 말하며, fully convolutional network를 제안한다. 특히, 매우 적은 수의 학습 데이터로도 학습이 가능하면서도, 더 정확한 segmentation 성능을 보여준다고 한다.
  U-Net 또한 FCN의 아디이어를 기반으로 효율적인 Segmentation Model을 제안한다. FCN에서는 CNN에 successive layers(==deconv Layers), 즉 pooling 연산자 대신 upsampling 연산자로 대체한 Layer를 추가함으로써 공간 정보를 추가하고, Segmentation의 성능을 끌어올릴 수 있었다. U-Net은 expanding Path(==upsampling 하는 부분)에 FCN보다 더 많은 conv filter들을 사용함으로써 higer resolution layers에 context information을 propagate 할 수 있다고 말한다. 또한 (FCN과는 달리) U-Net은 High resolution을 더 잘 처리하기 위해 overtile strategy를 사용한다. 이러한 전략을 선택해 더 큰 resolution에 대해서도 학습을 할 수 있도록 했다.

 

   Overtile strategy
   FCN 논문에서는, patch-wise learning에 대해서 부정적인 입장을 취한다. overlap 되는 부분 때문에 학습 성능이 저하될 수 있을 뿐 더러, patch-wise에서는 patch 별로 학습을 진행하기 때문에 연산량 또한 크다. 그래서 FCN 논문에서는 이미지를 자르지 않고 하나의 이미지로 학습을 진행한다는 내용이 언급되어 있다. 하지만 큰 이미지에 대해서는 어떻게 동작할 지에 대한 구체적인 방안은 언급이 없다. U-Net에서는 FCN이 잘못되었다고 말하지는 않으나, GPU Memory에 대해서 언급한다. High resolution을 처리하기 위해 HW적으로 한계가 있는 경우, 한 이미지로 처리하기 힘들 기 때문에 결국 이미지를 잘라서 처리한다.

   이 때문에 U-Net에서는 큰 이미지를 처리하기 위한 방법으로 Overlap-tile strategy를 제안 및 사용한다.

   기존 Patch 학습 방법에서는 Overlap 되는 부분에 대해서 그대로 학습을 진행했다면, U-Net에서는 Overlap되는 부분을 실제로 학습시키고 싶은 image의 내용을 mirrorring 해서 사용한다. 위 그림처럼 노란 그림에 대한 segmentation map을 얻고 싶을 때, 빨간 색 네모의 이미지를 집어 넣어서 사용하고, 노란 그림 바깥의 부분은 mirroring 해서 사용하는 것이다. (아래 모델 그림에서 resolution을 확인하면 572x572가 input으로 들어가서 388x388이 나오는 것을 확인할 수 있다) 이를 통해서 Overlap되는 부분 때문에 학습의 저하가 일어나는 것을 사전에 방지한다. 또한 U-Net에 vaild part라는 내용이 나오는데, 이 vaild part를 활용하는 것 또한 긍정적인 효과가 일어났다고 얘기한다. (padding을 사용하지 않고 mirrorring을 사용함으로써 일어나는 일..)

 

    Data Augmentation

    bio 쪽은 데이터가 부족하기 때문에 Data Augmentation이 굉장히 중요한 주제였던 것으로 보인다. 그래서 dataset 구축을 위해 elastic deformation을 진행했다고 한다. (Shift, Rotation, Gray value, Random elastic deformation)

 

Model 

U-Net의 구조는 크게 2가지 Path로 나누어져 있다.

 

1. Left side
  Left side는 contracting path, encoder 등으로 불리는 부분이다. 해당 부분은 일반적인 CNN과 결을 같이 하게 만들어 져 있다. U-Net은 padding을 하지 않는 컨셉을 유지하기 때문에, 3x3 conv를 반복해서 수행할 때마다 feature map의 resolution의 H, W가 각각 2씩 줄어드는 것을 확인할 수 있다.

  1. two 3x3 convolutions -> convolution을 두번 연속으로 시행한다고 함, unpadded convolutions을 한다(U-Net은 padding을 하지 않는 컨셉)
  2. Relu를 거친 후에, 2x2 max pooling 연산을 stride 2로 진행해서 downsampling을 한다. Downsampling을 할 때 channels의 개수를 2배로 늘린다.

 

2. Right side

  Right side는 expansive path, decoder 등으로 불리는 부분이다. feature map을 upsampling 하는 부분이다.

  1. 2x2 De-Conv filter를 통해서 upsampling을 진행 (FCN과 같이 Transposed Convolution 사용 or 보간법 사용)
      2x2 De-Conv(un-conv 2x2) == Transposed Convolution인데, 이를 수행하고 나면 feature map은 up-sampling 되고 feature map은 절반으로 설정하여 줄인다.

  2. Contracting path에서 같은 단계에 해당하는 부분에 mirroring 된 부분을 제거하고 현재 Layer의 feature map에 연결한다. (copy and crop)
  3. 2번까지 수행하면, 위 그림에서 하얀 박스와 파란 박스가 합쳐진 상태이다. 이 상태에서 3x3 Convolution (unpadded)를 진행하면서 또 Channel 수를 반으로 줄이도록 한다.

  4. 3x3 conv를 진행한 다음에 다시 De-Conv filter를 통해 upsampling을 진행하면서 Channel 수를 반으로 줄인다.

 

이러한 과정이 반복된다.

 

 

 

Reference

https://89douner.tistory.com/297