0. Abstract
기존에는 Deep Learning Network들을 성공적으로 훈련시키려면 수 천개의 annotated sample들이 필요했다.
그러나 U-Net의 경우 강력한 Data Augmentation 기법을 이용하여 few image 만으로도 효율적으로 모델을 훈련할 수 있도록 한다.
U-Net Architecutre는 크게 2가지로 구성되어 있다.
- Contracting Path: capture context
- Expansive Path: enable precise localization
Skip Architecture를 사용해서 Shallow, Fine Layer의 Appearance Information (Local, Where)과
Deep, Coarse Layer의 Semantic Information (Global, What)을 결합할 것이다.
Appearance Information
- Shallow, Fine Layer
- Local Information
- Location: Where
Semantic Information
- Deep, Coarse Layer
- Global Information
- Semantics: What
CNN에서 일반적으로 Classification Task를 수행할 때
Layer를 통과할수록 Low-level feature에서 High-level feature가 추출되어,
Hierarchical Information들을 Edge -> Texture -> Object 순서로 얻을 수 있다.
그렇기에 Shallow, fine Layer는 객체의 외관/윤곽과 같은 Edge Feature들인 Local Information (Where)을 잘 추출하고,
Deep, coarse Layer는 객체의 외관을 파악하기 힘들지만, feature들이 활성화된 부분을 보면 옷과 같이 하나의 의미를 갖는 정보들을 보여준다. 즉, 하나의 Context Feature인 Global Information (What)을 잘 추출한다.
아래의 그림을 보면 U-Net의 전반적인 구조를 이해할 수 있다. (U자 모양으로 생겨 U-Net이라 부른다)
1) Contracting Path을 통해 context information (Global, Semantics)을 구하고,
2) Contracting Path와 symmetric한 Expansive Path를 통해 precise localization (Locat, Spatial Details)을 수행하여 원본 이미지 크기와 동일하도록 맞춘 후 Segmentation을 진행한다.
FCN과 U-Net의 차이점은 '같은 층'에서 concatenation이 일어나는 것의 유무이다.
e.g. U-Net의 경우:
Contracting Path : $568 \times 568 \times 64$
Expansive Path : $392 \times 392 \times 128$
사이에는 맨 윗층에서 Concatenation이 일어나 shallow, fine layer의 apperance information을 deep, coarse layer의 semantic information에 더해주는 skip - connection이 일어난다고 볼 수 있다.
이러한 U-Net 모델은 ISBI cell tracking challenge 2015에 참가하여 좋은 퍼포먼스를 보여준 모델이기에 해당 challenge 대회에서 사용한 데이터셋인 "Transmitted light microscopy images"를 사용했다.
1. Introduction
첫 문단은 기존의 CNN 모델에 대해 간략하게 설명하고 있다.
AlexNet, VGGNet 에 대한 소개, 그리고
2015년 당시 소개된 최신 모델은 VGGNet or GoogLeNet 이었다.
기존의 CNN은 하나의 Image를 single class label로 반환하는 Classification Task에 사용되었으나,
Biomedical Image Processing 분야에서는 한 Image 내의 모든 Pixel에 대해 Classification을 수행하는 Segmentation Task를 수행하기 위해 다른 형식의 모델이 필요했다.
또한 해당 분야에서는 Biomedical Task를 수행하기 위한 training image들을 구하기 쉽지 않다는 문제가 있다.
이러한 문제들을 해결하기 위해 Patchwise 방법(Sliding Window)와 또 다른 Patchwise 방법(U-Net)이 나타났다.
Pixelwise Learning: FCN
Patchwise Learning: Sliding Window, U-Net
1.1) Sliding-Window Setup
Sliding Window는 Patchwise 방법이라고도 하는데, 하나의 이미지를 기준으로 Patch를 sliding 하면서 주어진 Patch의 center pixel을 분류하는 방식이다.
Patchwise Learning 방식은 아래와 같다.
1) 특정 size의 patch를 설정하고 이를 CNN에 입력으로 넣어준다.
2) 입력으로 들어간 patch에 대응되는 Input Image의 일부는 CNN에 의해 Classification Task가 수행된다.
3) 이때 특정 Class로 분류가 되었다면, 해당 patch 중앙에 위치한 pixel을 해당 class로 분류해 준다.
4) 이러한 과정을 Sliding Window 방식으로 1 ~ 3번을 반복한다.
다만, Patchwise Learning 방식에는 여러 문제점들이 존재한다.
1) Network는 각 Patch마다 독립적으로 CNN연산을 수행해야 하기 때문에 너무 많은 계산량을 요구하게 되어 속도가 느려진다.
2) 각 Patch들끼리 중복되는 부분이 있으나 독립적으로 CNN 연산이 수행되기에 이를 활용하지 못한다.
3) Patch size에 따른 localization accuracy & use of context 간의 trade-off
Patch Size가 커질 경우: 최종 Down-sampled map size가 동일하기에 Subsampling(Pooling) Layer들이 더 많이 요구된다.
즉, 기존 Patch에 대응되는 이미지에 Downsampling이 더 많이 적용되어 Local, Spatial Detail들을 덜 담게 된다.
이는 Context는 많이 보지만 Localization Accuracy가 떨어지는 결과를 낳는다.
또한 하나의 patch 안에 두 개의 class가 동시에 들어가 있을 확률이 커져 class label을 어떻게 부여할지에 대한 문제도 있고,
patch끼리 겹치는 부분이 커져 중복계산이 진행되는 경우가 많아진다.
Patch Size가 작아질 경우: 기존에 Patch에 대응되는 Input Image의 영역이 더 작아지기에 Network가
이전보다 더 작은 Context를 학습할 수 밖에 없다.
즉, 위의 Sliding Window 방식의 Patch Redundancy 문제를 해결하기 위해 U-Net에서는 Overlap되는 Patch가 없도록 학습한다.
그에 따라, U-Net의 Input은 Image 전체가 아니라 (Patch에 대응되는 Image의 일부 + Extrapolation된 부분)이 들어간다.
Output도 Image 전체에 해당하는 크기의 Segmentation Map이 아니라
Extrapolation 부분을 제외한 Patch에 대응되는 크기만의 Segmentation Map을 출력한다.
1.2) FCN & U-Net
그렇기에 논문의 저자들은 FCN(Fully-Convolutional Network)을 이용하는 방식을 사용했다.
따라서 기존의 FCN Architecture를 modify and extend 하여 few training images에 대해 학습하고 precise segmentation을 수행할 수 있도록 새로운 모델인 U-Net을 설계한다.
- Contracting Path: capture context
- Expansive Path: enable precise localization
아래의 그림을 보면 U-Net의 전반적인 구조를 이해할 수 있다. (U자 모양으로 생겨 U-Net이라 부른다)
1) Contracting Path을 통해 context information (Global, Semantics)을 구하고,
2) Contracting Path와 symmetric한 Expansive Path를 통해 precise localization (Locat, Spatial Details)을 수행하여 원본 이미지 크기와 동일하도록 맞춘 후 Segmentation을 진행한다.
3) 이러한 Expansive Path에서는 Upsamped Output과 Contracting Path에서의 High Resolution Features가 서로 combine되어 Skip-Connection을 형성한다.
4) Expansive Path에서는 기존 FCN보다 더 많은 Convolution Filter Channels가 사용되었다. 이는 Context Information을 High Resolution Layer로 더 효과적으로 전달할 수 있게 된다.
이러한 U-Net Architecture는 FC Layer들이 존재하지 않으며 Convolution Layer의 Feature Extraction 부분만 사용한다.
1.3) Overlap-tile Strategy
U-Net의 경우 Convolution Layer에서 Padding을 사용하지 않기에 Input Size보다 Output Size가 무조건 작다.
그렇다면 Input을 Image의 일부 영역 뿐만 아니라 Extrapolation을 통해 확장시킨 후,
Output은 그대로 Segmentation size로 고정시켜 학습을 진행한다면 Patch끼리 Overlap 되지 않는다.
즉, Patch에 대응되는 Image의 일부 + extrapolation된 부분의 size가 $572 \times 572 \times 1$이고,
실제 Segmentation 되는 부분은 extrapolation된 부분을 제외한 Patch에 대응된 Image의 size여서
Output Image의 Size가 $388 \times 388 \times 2$가 되는 것이다.
위의 그림에서 빨간색 영역과 주황색 영역이 각각 U-Net의 Input과 Output이 되는 것이다.
- Input: 실제 Patch에 대응되는 Input Image의 일부 + Extrapolation된 부분
- Output: 실제 Patch에 대응되는 Input Image의 일부
실제 Input은 Extrapolation된 부분까지 고려한 후에,
그에 맞춰 정사각형 형태로 Input Size를 설정한다.
해당 논문에서는 빨간색 Bounding Box와 실제 segmentation size에 해당하는 주황색 Bounding Box를 뺀 나머지 부분을 "missing context"라 한다.
그리고 mirroring 기법을 이용하여 missing context 부분을 extrapolation 하여 U-Net의 Input을 만들어준다.
이러한 방법을 Overlap-Tile Strategy라 부른다.
즉, 이를 통해 U-Net의 Input Size에 제한받지 않고 Segmentation Task를 수행할 수 있다.
어차피 Segmentation size는 고정되어 있기에,
Extrapolation 되는 부분만 늘리면 Input Size를 임의로 조절할 수 있다.
1.4) Challenges
논문에서는 2가지의 Challenge에 대해 언급했다.
첫 번째는, Biomedical segmentation task의 특성상 제한적인 training dataset을 지니고 있다.
따라서 excessive data augmention을 elastic deformation 방법을 이용하여 학습 데이터 수를 증가시켰다.
두 번째는, cell segmentation task의 경우 semantic segmentation으로는 동일한 class인 인접한 세포들을 분리하기가 쉽지 않다.
따라서, weighted loss를 고려하는 별도의 loss function을 사용하여 인접한 cell들 사이에 경계가 될 수 있는 background를 잘 segmentation할 수 있도록 했다.
Instance Segmentation: 같은 Class의 Instance를 구별하는 segmentation
Semantic Segmentation: 같은 Class의 Instance를 구별하지 않는 segmentation
2. Network Architecture
Contracting Path
Contracting path는 전형적인 CNN의 Feature Extractor 부분을 따른다. (FC Layer가 제거된 부분)
1) 두 번의 $3 \times 3$ Convolution Layer $(3, 1, 0)$, ReLU Activation
Convolution Layer Parameter: (kernel_size, stride, padding)
$(3, 1, 1)$: Kernel Size 유지
$(3, 2, 1)$: Kernel Size를 절반으로 Down-Sampling
$(3, 1, 0)$: Kernel Size - 2 만큼 Down-Sampling
2) $2 \times 2$ Max-Pooling Layer (stride=2), Down-Sampling시 Feature Channel $\times 2$
Expansive Path
Expansive Path는 Contracting Path와 대칭적인 구조를 띄고 있다.
1) Down-Sampling을 수행하는 Max-Pooling Layer 대신 Up-Convolution Layer로 변환
실제 코드를 찾아보니 Up-convolution에 사용되는 연산은 $2 \times 2$ Transposed Convolution $(2, 2, 0)$이다.
Up-Sampling시 Feature Channel $\div 2$
Transposed Convolution Layer Parameter: (kernel_size, stride, padding)
$(3, 1, 1)$: Kernel Size 유지
$(4, 2, 1)$: Kernel Size를 2배로으로 Up-Sampling
$(2, 2, 0)$: Kernel Size를 2배로 Up-Sampling w/ No Padding
2) Copy and Crop: Contraction Path에서 추출된 Feature Map을 Crop하여 Concat 한다.
앞에서 Mirroring을 통해 extrapolation된 부분만 Crop을 통해 100% 제거되는 건 아니다.
Extrapolation된 부분 뿐만 아니라 기존의 Input의 일부도 같이 고려된 부분도 Crop될 수 있다.
따라서 Expansive Path의 Feature Map size와 맞춰주기 위해 Crop 한다.
e.g. contracting path에 있는 $136 \times 136$ Feature Map을 Crop하여 $104 \times 104$ Feature Map을 만들고
이를 Up-Sampling된 $104 \times 104 \times 256$ Expansive Path의 Feature Map과 Concat 한다.
3) 두 번의 $3 \times 3$ Convolution Layer $(3, 1, 0)$, ReLU Activation
4) Final Layer: $1 \times 1$ Convolution을 이용하여 $64$개의 Channel을
desired number of classes인 $2$개의 class로 segmentation 하기 위해 $2$개의 Channel을 가지도록
$388 \times 388 \times 64$ Feature Map을 $388 \times 388 \times 2$의 Output Segmentation Map으로 만든다.
Desired number of classes = $2$인 이유는
Bakground(배경) Class와 Foreground Class가 존재하기에 2개이기 때문이다.
따라서, U-Net은 총 $23$개의 Convolutional Layer를 가진다.
$3 \times 3$ Convolutional Layer $(3, 1, 0)$: $18$ (blue arrows)
$2 \times 2$ Transposed Convolutional Layer $(2, 2, 0)$: $4$ (green arrows)
$1 \times 1$ Convolutional Layer $(3, 1, 0)$: $1$ (yellow arrows)
Overlap-tile strategy가 잘 적용되기 위해, $3 \times 3$ Convolutional Layer의 input으로 들어오는
size $(x, y)$의 $x, y$를 모두 짝수로 설정해야 $2 \times 2$ Max - Pooling이 잘 적용된다.
3. Training
Input Image와 Output Segmentation Map을 가지는 U-Net은 SGD를 기반으로 Training된다.
본 논문에서 사용하는 Deep Learning Framework는 Caffe이다.
또한 U-Net이 사용하는 Convolution Layer의 경우 Unpadded Convolution이므로 무조건 Input size가 Output size보다 크다.
이때, Training시 Patch size를 매우 크게 하여 (Large Input tiles) Batch를 single image로 줄인다.
이렇게 할 경우 mini-batch SGD에서 Gradient가 수렴하는 방향으로 가지 않을 수 있기에
Momemtum을 $0.99$라는 매우 큰 값을 주어 이전에 보았던 수 많은 training sample들이
현재 optimization step에서 gradient를 update하는데 영향을 준다.
이는 일종의 mini-batch에서 오는 문제점을 해결해준다고 볼 수 있다.
정리하자면 Training의 Parameter는 다음과 같다.
Optimizer:
- Stochastic Gradient Descent (SGD)
- Momentum = $0.99$
Deep Learning Framework: Caffe
Batch: A single image
3.1) Loss Function
U-Net에서 사용하는 Loss Function은 아래와 같은 Energy Function이다.
$E = \sum_{x \in \Omega} w(x) \log(p_{\ell(x)}(x))$
Energy Function은 final feature map에다 픽셀 단위로 softmax를 수행한
pixel-wise softmax $P_{k}(x)$와 cross-entropy를 결합하여 Loss Function을 적용했다.
(1) Softmax
$p_k(x) = \frac{\exp(a_k(x))}{\sum_{k'=1}^K \exp(a_{k'}(x))}$
$K$: Number of classes
$k$: $k^{th}$ Class
$x \in \Omega \: with \: \Omega \subset Z^{2} $: $x$(=logit)의 pixel position $(i, j)$
$a_{k}(x)$: Softmax input ($k^{th}$ channel의 pixel position $x$의 activation)
주의할 점은 activation의 의미가 final layer에서의 activation이 아니라
final layer의 직전 hidden layer에서 출력된 $k^{th}$ channel의 pixel $x$의 output activation이라는 점이다.
Classification softmax와 조금 다른 점은 softmax의 입력 변수인 $x$(=logit)가 2D Dimension 정보 $(i, j)$를 implicit하게 담고 있다는 점이다.
그 이유가 segmentation task에서는 pixel의 좌표 $(i, j)$마다 softmax 값을 구하고 classification (per pixel) 해줘야 하기 때문이다.
즉, logit에 해당하는 $x$는 특정 pixel 좌표의 정보를 represent 한다고 할 수 있다.
여기서 중요한 점은 Softmax Layer를 하나의 Feature Map에 있는 각 pixel에 적용하는 것이 아니라,
동일한 pixel position $x$에 각 $1 \sim K^{th}$의 Channel마다 Softmax Layer를 적용하는 것이다.
그래야 총 $K$개의 Class에 대해 Softmax를 하는 $\sum_{k'=1}^{K}$이 성립하여 어느 Channel에 속하는지 알 수 있다.
따라서, pixel position $x$가 background(배경)와 foreground(세포) 중 어느 Class에 속하는지
각 $k$의 값에 따른 $p_{k}(x)$의 값을 비교하여 판단할 수 있다.
$p_{k}(x) \approx 1$: pixel position $x$가 $k^{th}$ channel(class)에서 maximum activation $a_{k}(x)$을 가진다.
$p_{k}(x) \approx 0$: pixel position $x$가 그 외의 $all \: other \: k$에서는 $k^{th}$ channel(class)에 속하지 않기에 0에 가까워진다.
(2) Cross-Entropy Loss
논문에서 소개된 Segmentation Task는 Class가 2개이므로 (Background, Foreground) Binary Cross-Entropy Loss를 사용한다.
Cross - Entropy Loss = $- \sum_{N} (Real \: Value) \times log(Prediction)$
여기서 Real Value는 One-Hot Encoded Vector이다.
$E = \sum_{x \in \Omega} w(x) \log(p_{\ell(x)}(x))$
$l: \Omega \rightarrow {1, ..., K}$: pixel position $x$가 갖는 true label $l(x)$
$w: \Omega \rightarrow $: weight map $w$
$k$자리에 $l(x)$을 넣어서 true label $l(x)$을 가지는 예측값을 Softmax Activation $p_{l(x)}(x)$ 으로 처리한다.
(3) Weight Map
$w(x) = w_c(x) + w_0 \cdot \exp\left(-\frac{(d_1(x) + d_2(x))^2}{2\sigma^2}\right)$
일단 Instance Segmentation task를 위해 같은 class(=cell)들을 구분해줄 수 있는 경계 (boundary)에 해당하는
pixel을 더 잘 분류해줄 수 있도록 boundary에 속한 pixel들의 가중치를 더 강하게 주어 더 큰 loss 값을 출력해
해당 pixel(=boundary) 들을 더 잘 분류할 수 있도록 학습시킨다.
$W_{c}: \Omega \rightarrow R$: Class Frequencies를 Balance하기 위해 사용
Training data에 있는 class(background, foreground)의 빈도수를 balance하기 위해 weight를 사용한다.
$W_{c}(x) = \frac {1} {frequency}$
$frequency = \frac {Pixel \: Count} {Total \: Number \: of \: Pixels}$
Class의 빈도수가 높을수록 $W_{c}$가 낮게 나오기에 Class frequency를 balance할 수 있다.
cell을 구분 짓는 boundary 영역에 위치한 pixel의 $d_{1} + d_{2}$값이 작게 나온다.
즉, $d_{1} + d_{2}$값이 작게 나올수록 높은 가중치가 부여된다고 볼 수 있다.
$w: \Omega \rightarrow R$: Weight Map
$w_{0}$: Constant (상수)
$d_{1}: \Omega \rightarrow R$: $x$ position에 있는 pixel에서 가장 가까운 cell border까지의 distance
$d_{2}: \Omega \rightarrow R$: $x$ position에 있는 pixel에서 두 번째로 가까운 cell border까지의 distance
이때, 측정하는 Distance인 $d_{1}, d_{2}$는 pixel을 중심으로 동심원을 그리다가 cell의 경계에 접하는 가장 첫번째 순간에 해당하는 point와 두 번째 순간에 해당하는 point까지의 반지름이라 생각하면 된다.
즉, Pixel $x$ ~ Cell의 경계까지의 거리 $= d$ 이라고 볼 수 있다.
본 논문에서는 $w_{0} = 10$, $\sigma \approx 5 \: pixels$의 값을 사용했다.
Q. 왜 두 Cell들 간의 Distance를 고려하는데 $(d_{1}(x) + d_{2}(x))^{2}$으로 합의 제곱을 구하는가?
한 pixel에서 두 Cell들 간의 거리의 합을 기준으로 둔다면, $d_{1}(x) +d_{2}(x)$가 더 커도 실제 두 cell들 간의 거리가 더 작을 수 있지 않나?
만약 그렇다면 두 cell간의 거리를 구하는 기준이 잘못된 것이 아닌가?
A. $d_{1}(x) + d_{2}(x)$가 최소가 될 때 예외 없이 모든 case에 대해 두 cell들 간의 거리가 최소가 된다고 보자.
그렇게 된다면 하나의 pixel $x$에 대해 $d_{1}(x) + d_{2}(x)$가 최소가 되는 순간이 존재하여 두 cell간의 거리가 최소가 된다.
Q. 왜 굳이 Gaussian Distribution을 따르는가?
$\exp\left(-\frac{(d_1(x) + d_2(x))^2}{2\sigma^2}\right)$에서 2D Gaussian Distribution이라면 식이
$\exp\left(-\frac{d_1(x)^{2} + d_2(x)^{2}}{2\sigma^2}\right)$으로 중심이 $(d_{1}, d_{2})$이 되어야 하는거 아닌가?
A. 만약 $d_{1}(x) + d_{2}(x)$이 하나의 distance value로 distance가 최소일 때 weight을 크게 한다면 그냥 2차원 평면상에서 해석하는 것이 아니라 scalar value로 해석해서 1D Gaussian Distribution을 적용해도 되지 않나?
3.2) Weight Initialization
U-Net의 경우 weight initialization은 He Normal Initialization을 사용한다.
$W \sim N(0, \text{Var}(W))$
$\text{Std}(W) = \sqrt{\frac{2}{n_{\text{in}}}}$
$n_{in}$: 이전 Layer의 Input Node 수
그렇기에 Standard Deviation은 $\sqrt{\frac{2}{N}}$이고,
$N = k^{2}C =$ Convolutional Filter Size $\times$ Convolutional Filter Channel
3.3) Data Augmentation
Data Augmentation은 training sample들이 매우 적을 때 desired invariance와 robustness properties를 네트워크에 학습시키기 위해 사용하는 필수적인 기법이다.
U-Net에서 언급하고 있는 Data Augmentation Method는 4가지가 있다.
1) Shift
2) Rotation
3) Gray Value
4) Random elastic deformation
Elastic Deformation
pixel이 random하게 다른 방향으로 뒤틀리도록 변형하는 방식
이는 Cell들과 유사한 형상이므로 Elastic Deformation을 하여 실제로 있을 법한 데이터로 증강시킴
저자들은 다음과 같은 3-step을 통해 Data Augmentation을 수행했다.
1) coarse $3 \times 3$ grid에 random displacement vector를 이용해 smooth deformation을 수행한다. displacement는 10개 pixel들의 표준편차를 따르는 Gaussian Distribution에서 임의로 뽑은 값으로 수행한다.
2) Bicubic Interpolation을 이용해 pixel 단위로 displacement를 계산한다.
3) Contracting Path의 끝에 있는 Drop-out Layer가 더욱 implicit한 Data Augmentation을 수행한다.
4. Experiments
U-Net을 이용한 3가지 서로 다른 Segmentation Task에 적용한 Experiment에 대해 설명한다.
1) The segmentation of neuronal structures in electron microscopic(EM) recordings
2) Cell segmentation task in light microscopic images
2.1) Dataset: PhC-U373
2.2) Dataset: IDC-HeLa
Evaluation Metric는 총 3개가 있다: Warping Error, Rand Error, Pixel Error
pixel error는 prediction과 ground truth에 Pixel 마다 할당된 class들끼리 불일치하는 pixel들의 개수가 전체 pixel들을 기준으로 어느정도의 비율인지 알려주는 metric이다.
Warping Error와 Rand Error는 Digital topology error와 관련이 있는듯 하다.
해당 부분은 위상수학에 대한 이해를 기반으로 해석해야 하는데, 시간 부족으로 인해 해당 부분까지는 살펴보지 못했다.
https://imagej.net/plugins/tws/topology-preserving-warping-error
4. Conclusion
U-Net Architecture는 biomedical segmentation application에 대해 매우 우수한 성능을 보인다.
elastic deformation을 이용한 data augmentation 기법을 통해 매우 적은 수의 annotated image를 가지고 합리적인 training time을 가지고 학습이 가능하다.
U-Net은 비단 biomedical segmentation task뿐만 아니라 다양한 Segmentation task에도 활용 가능하다.
6. References
https://89douner.tistory.com/297#recentComments
https://youtu.be/NhdzGfB1q74?si=tdyGxVpisPhWhDVl
https://velog.io/@minkyu4506/%EB%85%BC%EB%AC%B8-%EB%A6%AC%EB%B7%B0-U-Net-%EB%A6%AC%EB%B7%B0
https://youtu.be/O_7mR4H9WLk?si=-caym5EP6N1rzdht
https://youtu.be/evPZI9B2LvQ?si=tRFE3KqydxwdqcOI