대표적인 Semantic segmentation 모델인 U-Net에 대해 알아보자!
원문 링크 : https://arxiv.org/abs/1505.04597
먼저 알아두면 좋은 것 - FCN : https://arxiv.org/abs/1411.4038
시작하기 전에 먼저 Semantic Segmentation이 무엇인지 알아보자.
Semantic Segmentation (의미론적 분할) 이란 영상의 모든 픽셀을 분류하여 클래스별로 분류하는 작업을 말한다.
위 그림을 예시로 살펴보면 개와 고양이, 그리고 배경으로 이루어진 이미지가 주어졌을 때 이를 네트워크에 통과시킨 결과는 각 픽셀이 어느 class에 속하는지 나타내는 segmentation map이다.
이는 곧 모든 픽셀마다 classification을 수행하는 것과 같다.
U-Net Introduction
U-Net이 나올 무렵 Image recognition task에서는 CNN이 SOTA를 찍고 있었다.
이 때 CNN의 주된 사용 목적은 Classification이었다.
앞서 설명했듯 segmentation은 output의 localization이 필요하고 이는 즉 픽셀마다 class label이 할당되어야 함을 뜻한다.
그러나 classification을 위한 CNN 구조의 경우 spatial information이 중요한 segmentation task에 적용하기에는 무리가 있다.
(pooling 등 연산을 거치면 spatial resolution이 줄어들고 fully connected layer를 통과시키면 공간적 정보를 잃게 된다!)
Segmentation을 수행하기 위한 Sliding window setup과 같은 시도들도 있었는데,
- 이미지에 임의의 크기의 "window"를 두고
- 이미지 전체에 대해 patch를 만든 후
- patch마다 network를 적용해 class label을 예측하도록 하였다. 이 때 픽셀 주변의 local region 정보를 이용한다.
그러나 이 방법의 경우 각 patch마다 네트워크를 적용하므로 속도가 느리고 겹치는 부분이 많아 연산량이 많았다.
또한 patch 크기가 커질 수록 max pooling 수가 늘어나 Localization accuracy가 감소하나 반대로 patch 크기를 작게 하면 네트워크가 좁은 영역의 context밖에 볼 수 없어 tradeoff 문제가 있었다.
따라서 Fully Convolutional Network 구조를 차용한 U-Net이 제안되었다.
U-Net은 적은 수의 데이터로도 정확한 segmentation이 가능하다는 장점이 있다.
Main Idea of U-Net
Contracting Path & Expanding Path
위의 구조를 살펴보면, convolution과 pooling이 수행되는 부분과 convolution과 upconvolution이 수행되는 부분으로 나뉘는 것을 확인할 수 있다. 앞의 부분을 contracting path, 뒤의 부분을 expanding path라고 부른다.
Contracting path에서는 feature를 추출하고, Expanding path는 feature map을 다시 input resolution과 가깝게 upsampling하는 역할을 한다. 이 때 contracting path에서의 feature map이 expanding path에서의 feature map과 concat되는 것을 볼 수 있는데 high resolution feature map을 upsampled된 feature map과 합쳐 localization을 더 잘 하기 위함이다.
Large number of channels in the upsampling part
upsampling part가 많은 feature channel을 갖기 때문에 높은 resolution layer로 context 정보를 전파할 수 있다.
No Fully Connected Layers
오직 convolution layer만 존재한다. (공간 정보를 잃지 않기 위해서!)
Tiling Strategy & Data Augmentation
이미지의 경계에 있는 픽셀은 mirroring을 통한 extrapolation으로 segmenation하였다. 또한 U-Net은 Biomedical task에 집중하는데, 문제 특성상 많은 데이터를 확보하기 어렵기 때문에 data augmentation을 통해 deformation-invariant한 학습을 수행하도록 하였다.
Network Architecture
Contracting Path
일반적인 CNN 구조와 같다. 한 번의 downsampling(pooling)마다 channel 수가 2배로 증가한다.
2번의 3x3 convolution을 거친 후 한 번의 pooling (2x2, stride 2)을 반복한다.
Expanding Path
contracting path 구조와 비슷하나 pooling을 upconvolution 연산으로 대체한다.
한 번의 upsampling마다 channel 수가 1/2로 감소한다. 마찬가지로 2번의 3x3 convolution 후 한 번의 upsampling을 한다.
Concat
각 path에서 각 단계마다 feature map을 concat해 준다.
이 때 expanding path에서의 feature map size가 더 작기 때문에 contracting path에서의 feature map을 crop한 뒤 concat 해 준다. → 이 때 채널 개수가 다시 2배로 늘어나게 된다.
마지막 layer에서는 1x1 convolution을 수행해 feature vector를 class로 mapping한다.
Pytorch implementation (수정중)
https://github.com/Minjee-Seo/Pytorch/blob/main/Implementation/UNet.py