목차
1. U-Net 배경
2. U-net 구조
2.1 개요
2.2 아키텍쳐
2.3 Overlap-tile starategy
3. Training Details
3.1 Optimization
3.2 Loss Function
3.3 Data Augmentation
4. Performance
5. Reference
1. U-Net 배경
본 논문의 과제는 세포 이미지(그림 1)에서 세포막과 배경, 서로 맞닿은 세포막(Touching objects of the same class)을 구분하는 세그멘테이션이다. 세그멘테이션을 위해서는 두 가지 조건이 있다. 먼저 위치 검출(Localization) 성능이 좋아야 한다. 이미지 분류(Image Recognition) 과제에서는 이미지에 어떤 객체가 있는지 분류만 하면 된다. 세그멘테이션에서는 분류하고자 하는 클래스별로 객체가 어디에 위치해 있는지도 정확히 알아야 한다. 두번째로 Deep CNN기반 세그멘테이션 모델을 학습시키려면 데이터가 많아야 한다. Semantic Segmentation 모델은 픽셀마다 레이블된 이미지로 학습한다. 그런데 바이오 분야에서는 픽셀마다 레이블된 이미지를 구하기 어렵다. 여러 가지 이유가 있겠지만, 일상적인 이미지는 누구나 레이블링할 수 있는 반면 바이오 이미지를 레이블링할 수 있는 인력은 일정 수준의 전문 지식을 갖춰야 하기 때문일 것이다.
U-Net 이전에 바이오 분야의 이미지 세그멘테이션은 패치 단위 학습을 시도했다. 패치 단위 학습은 이미지 영역 하나로 픽셀 하나를 분류한다. 그림 2는 Ciresan 외 7명이 제시한 패치 단위 학습 방식이다. Original Image를 보면 픽셀 p를 중심으로 가로, 세로 길이가 w인 윈도우가 있다. 이 윈도우를 '패치'라 한다. 패치를 신경망(DNN)에 입력해 픽셀 p를 분류한다. 동일한 방법으로 모든 픽셀을 분류한다. 이때, 패치의 크기는 w로 일정하다.
패치 단위 학습은 세그멘테이션의 두 가지 조건을 만족할 수 있다. 먼저 이미지 데이터가 부족해도 괜찮다. 이미지의 픽셀 수만큼 패치를 만들 수 있기 때문이다. 예를 들어 50x50 이미지 한 장에서 총 2500(50x50)개 패치 데이터를 얻을 수 있다. 또한, 객체의 위치를 파악할 수 있다. 원래 전체 이미지 세그멘테이션 결과는 픽셀마다 분류한 클래스를 원래 위치에 따라 조합했기 때문이다.
하지만 Patch 단위의 학습 방식도 완벽하지는 않다. 일단 속도가 굉장히 느리다는 단점이 있다. 픽셀마다 분류를 해야 하기 때문에, 최종 세그멘테이션 결과를 얻으려면 많은 시간이 소요된다. 특히 Patch 데이터 중에는 원래 이미지에서 겹치는 부분이 많기 때문에 똑같은 부분을 반복적으로 연산하기 때문에 비효율적이다. 그리고 윈도우 크기에 따라서 객체 위치 검출(Localization)과 분류를 위한 맥락 정보(Context) 수준 사이에 트레이드오프(trade-off) 문제도 있다. 윈도우를 크게 하면 max-pooling을 더 많이 이용해서 위치 정보가 소실된다. 이는 세그멘테이션 결과 실제와 비슷한 모양으로 세그멘테이션해도, 정확한 위치에 세그멘테이션하지 못하는 문제를 말한다. 윈도우를 작게 하면 한 픽셀을 예측하는 데 맥락 정보(Context)가 그만큼 적어진다. 그만큼 픽셀 마다 분류의 정확도가 낮아질 수 있다.
본 논문은 이런 한계를 지적하며, U-Net 모델을 제시한다. U-Net은 FCN과 같은 end-to-end 모델이다. 한 번의 모델 연산으로 이미지 전체를 세그멘테이션할 수 있다. 즉 패치를 이용하지 않는다. 따라서 윈도우 크기에 따른 Localization과 Classification 사이의 트레이드오프 문제가 없다. 바이오 이미지 데이터가 적은 문제는 데이터 증강(Data Augmentation)으로 해결한다.
2. U-Net 구조
2.1 아키텍쳐
U-Net은 그림 3처럼 그래프로 그렸을 때 U자 모양으로 생겨서 붙여진 이름이다. 그림 3을 반으로 갈랐을 때 왼쪽 부분, 즉 입력된 데이터(input image tile)가 빨간색 화살표를 따라 아래로 떨어지는 구간을 Contracting Path라고 한다. 반대로 그림의 오른쪽에서 초록색 화살표를 따라 위로 올라가는 구간을 Expanding Path라고 한다. Contracting Path와 Expanding Path는 서로 거울에 비춘 것처럼 좌우가 대칭인 모양이다. 두 Path의 기능도 정반대인데, 말 그대로 Contracting Path는 이미지를 수축하는 구간, Expanding Path는 확대하는 구간이다.
Contracting & Expanding Path의 역할
Contracting Path는 특징을 추출(feature extractor)한다. Contracting Path에 사용되는 층(layer)은 필터 사이즈가 3x3, 스트라이드가 1인 convolution layer(conv 3x3)와 윈도우 사이즈가 2x2, 스트라이드가 2인 max pooling layer(max pool 2x2)다. conv 3x3은 특징맵의 패턴을 잡아내고, max pool 2x2는 특징맵의 정보를 최대한 보존하면서 크기를 축소(1/2x downsample)한다. Expanding Path는 축소된 특징맵의 크기를 키우고, 마지막에는 픽셀별로 클래스를 예측하는 역할을 한다. Expanding Path는 up-conv 2x2와 Contracting Path에도 사용된 conv 3x3, 마지막에는 필터 사이즈가 1x1이고 스트라이드가 1인 convolution layer(conv 1x1)로 구성된다. up-conv 2x2는 축소된 특징맵의 사이즈를 키우고(2x upsample), conv 3x3 는 Contracting Path에서와 마찬가지로 upsample된 특징맵의 패턴을 잡아낸다. 중요한 점은 up-conv 2x2에서 단순한 논리로 크기만 키우는 것이 아니라(ex. bilinear-interpolation), 더 좋은 결과를 내게끔 필터를 학습한다는 것이다. 마지막 conv 1x1에서는 세포막과 배경 2개 클래스로 분류하기 위해 64채널 특징맵을 2채널 특징맵으로 매핑한다.
Details
더 구체적으로 살펴보자.
활성화 함수 : 매 conv 3x3과 up-conv 2x2에서 나온 특징맵은 ReLU층을 통과한다.
특징맵의 채널 수 : max pool 2x2에서 downsampling된 후에는 conv 3x3의 필터 수가 2배로 늘어난다. 이는 곧 특징맵의 채널 수가 2배로 늘어난다는 말과 동일하다. 반대로 Expanding Path에서는 up-conv 2x2로 바로 특징맵 채널 수를 조절할 수 있다. upsampling할 때마다 이전보다 필터 수를 1/2로 줄여 채널 수를 차차 줄이게 된다.
Input과 Output 사이즈 차이 : conv 3x3층은 padding을 하지 않기 때문에, 특징맵의 사이즈가 가로, 세로 2줄씩 줄어든다. Contracting Path 뿐만 아니라 Expanding Path에서도 을 통과하면 사이즈가 2씩 줄어들기 때문에 Expanding Path의 특징맵은 대응되는 Contracting Path의 특징맵보다 작다. 그리고 Upsampling할 때마다 대응되는 Expanding Path와 Contracting Path의 특징맵의 크기 차이는 더 커진다. segmentation이 제대로 되려면 Input 이미지의 크기와 Output segmentation map의 크기가 동일해야 한다. 본 논문에서는 Output이 Input보다 작아지는 문제를 보완하기 위해 mirroring 방법을 사용한다. 이 방법에 대해서는 2.3에서 설명하겠다.
copy and crop : U-Net에서는 DeconvNet[3]과 마찬가지로 Upsampling할 때 Downsampling 정보를 이용한다. Deconvnet에서는 max-pooling할 때 선택된 값의 위치정보를 switch variable에 저장했다가, max-pooling과 대응되는 max-unpooling에서 값을 매핑할 위치를 선택할 때 이용한다.[3] U-Net에서는 up-conv 2x2 결과를 대응되는 Contracting Path의 downsample 단계의 마지막 conv 3x3의 결과 특징맵을 이어붙인다(concatenation). 위에 설명한 바와 같이 이어붙일 두 특징맵의 크기가 다르기 때문에, conv 3x3의 특징맵을 up-conv 2x2의 특징맵 사이즈에 맞춰 가장자리를 잘라낸다(crop). 본 논문에서는 copy and crop의 기능을 명확히 설명하지 않았지만, 순전파(forward propagation)에서 입력 데이터에 가까운 정보를 추가하는 역할과 역전파(backward propagation)에서 기울기(미분값)를 바로 Contracting Path로 보내 기울기 소실 문제를 완화하는 역할을 한다고 본다.
2.2 Overlap-tile & Mirroring strategy
Overlap-tile과 mirroring을 이용하면 큰 이미지도 segmentation할 수 있고, Output 사이즈가 작아지는 문제도 해결할 수 있다. 큰 이미지를 한 번에 segmentation 하기에는 GPU 메모리에 부담이 될 수 있다. 메모리 문제를 해결하기 위해서 Overlap-tile을 이용한다. tiling은 말 그대로 이미지를 화장실 타일처럼 쪼개서 U-Net의 Input으로 사용하는 것이다(그림3에서 input image tile이라고 쓰인 이유다). 그림 4에서는 노란색 박스가 이미지 타일 하나다. 이미지 하나를 6개 타일로 나눠서 예측하는 것으로 보인다. 파란색 박스는 노란색 박스를 segmentation하기 위한 input image tile이다.
앞서 U-Net이 Input 사이즈보다 Output 사이즈가 작다는 것을 확인했다. 파란색 박스가 노란색 박스의 크기 차이도 같은 맥락이다. Input과 Output의 크기 차이는 tiling 방법에서 input image tile끼리 겹치는(overlap) 현상으로 이어진다. 그림 4의 노란 타일과 바로 오른쪽 타일을 예측하기 위한 input image tile을 상상해보면 상당부분 겹치는 것을 알 수 있을 것이다.
[1]에서 픽셀마다 patch를 사용해서 같은 부분은 반복적으로 연산하는 비효율에 비하면, overlap-tile의 겹치는 정도는 굉장한 개선이라고 볼 수 있다. 문제는 그림 4의 파란색 박스처럼 원래 이미지보다 더 큰 tile을 이용해 segmentation하기 때문에, 빈 부분이 있다는 것이다. U-Net에서는 이 문제를 mirroring extrapolate(보외법, 외삽법)으로 해결한다. 그림 4의 왼쪽 이미지와 같이 input image tile이 삐져나오는 만큼 원래 이미지를 거울에 비추듯이 복사해 빈 부분을 채우는 방법이다.
2.3 FCN vs DeconvNet vs U-Net
FCN, DeconvNet, U-Net을 비교해보면, U-Net의 특징을 이해하기 좋다. U-Net의 기틀이 된 아이디어는 "Fully Convolutional Networks"[2]다. U-Net과 비슷한 시기에 노현우, 홍성훈, 한보형 씨가 발표한 "Deconvolution Networks"[3]도 FCN[2] 기반의 네트워크를 제시한다. 잠시 FCN과 DeconvNet을 복습해보자. FCN은 end-to-end 모델이다. 즉, 한번의 모델 연산으로 이미지에 대한 세그멘테이션 결과를 낸다. DeconvNet은 크게 3단계로 나뉜다. Edge Box를 이용해 이미지를 50개 관심 영역으로 나눈다. 관심 영역을 세그멘테이션한다. 관심 영역 세그멘테이션 결과를 합쳐 이미지 세그멘테이션 결과를 얻는다.
U-Net과 DeconvNet은 비슷한 구조를 가지고 있다. 다운샘플링(Downsampling) 네트워크와 업샘플링(Up sampling) 네트워크가 서로 대칭이라는 것이다. 또한 다운샘플링 뿐만 아니라 업샘플링 네트워크도 학습할 수 있다는 것이 공통점이다. 참고로 FCN의 업샘플링도 학습 가능하다는 것을 명시했지만, 실제 구현할 때는 쌍선형 보간법(bilinear-interpolation)만 이용해 업샘플한다.
하지만 업샘플링 네트워크를 구성하는 층이 다르다. DeconvNet에서는 unpooling 층와 upconvolution 층([3]에서 deconvolution 층)로 구성되고, U-Net에서는 upconvolution layer와 convolution layer로 구성된다.
또, DeconvNet과는 다르게 U-Net에서는 Expanding Path에서 매 unpooling 결과에 대응하는 Contracting Path의 특징 맵(Feature Map)을 붙여서(Concatenate) 사용한다. 그리고 DeconvNet은 Output의 크기가 Input의 크기와 동일한데 반해, U-Net은 Output의 크기가 Input보다 작아진다는 차이가 있다.
| FCN | DeconvNet | U-Net | |
| 대칭 구조 | X | O | O |
| 업샘플링 네트워크 학습 | △ | O | O |
| 업샘플링 네트워크 구성 | upconvolution 층 |
upconvolution, max-unpooling 층 |
upconvolution, convolution 층 |
| upconvolution의 역할 | bilinear-interpolation | 특징맵 Densify | 업샘플링 |
| Skip Connection | Summation | X | Concatenation |
| Input & Output size | Input = Output | Input = Output | Input > Output |
3. Training Details
이번 장에서는 U-Net을 훈련할 때 사용하는 Optimization, Loss function, Data Augmentation을 소개한다.
3.1 Optimization
최적화 방법으로는 SGD를 이용한다. GPU 메모리를 최대한 활용하기 위해서 input tile의 크기와 배치 사이즈를 크게해 배치 하나를 하나의 이미지로 구성한다. 모멘텀은 0.99로 설정하여 이전 이미지들의 상당수가 현재 단계의 가중치 업데이트에 반영되게끔 한다.
Q : 실제 구현할 때 타일 크기, 타일 개수, 배치 사이즈는 어떻게 정해지는가.
3.2 Loss Function
- x : 픽셀 위치, x ∈ Ω with Ω ⊂ Z²
- ak(x) : 특징맵의 k 채널, x 위치의 활성화 값
- K : 클래스 수
- pk(x) : 소프트맥스 함수. 가장 큰 활성화 값(ak(x))을 갖는 채널 k에 대해서는 pk(x)가 1에 근사하고, 나머지 채널에 대해서는 pk(x)가 0에 가깝다.
마지막 conv 1x1을 통과한 score map은 클래스 수(K)만큼의 채널 수를 갖는다. 세포와 배경을 구분하는 모델에서는 채널 수가 2개다. score map의 각 채널은 픽셀별로 특정 클래스에 속할 점수 값이다. 두 채널의 동일한 위치의 픽셀로 소프트 맥스 함수값을 구하면 가장 큰 점수를 갖는 채널의 함수값이 1에 근사하고 나머지는 0에 근사한 값을 갖는다.
- ℓ : Ω -> {1, ... , K}
- ℓ(x) : x 픽셀의 참값(label). k채널의 segmentation 맵에서는 k가 클래스를 의미하기 때문에 ℓ(x) = k다.
- pℓ(x)(x) : segmenation map에서 x 픽셀의 참값 채널(ℓ(x))의 소프트맥스 함수값.
- w : Ω -> R
에너지 함수는 픽셀별로 cross-entropy를 구해 합한다. 그런데 일반적인 cross-entropy가 실제값과 예측확률의 로그값을 곱하는 것과는 달리 w(x)라는 새로운 weight값을 이용한다. 이 weight값을 구하는 공식은 아래 식3과 같다.
- wc : Ω -> R. 클래스 빈도 균형을 맞추기 위한 weight map
- d₁ : Ω -> R. x픽셀과 가장 가까운 셀의 경계까지의 거리
- d₂ : Ω -> R. x픽셀과 두번째로 가까운 셀의 경계까지의 거리
- U-Net 논문에서는 w0 = 10, σ는 5픽셀에 근사한다.
Q : 실제 마스크는 instance wise segmentation map인 것인가?
weight map, w(x)의 역할은 두 가지다. 설명에 앞서 식 3에서 wc(x)를 균형 가중치항, wc(x)를 제외한 나머지를 경계 가중치항이라 하겠다. 하나는 실제 레이블된 segmentation mask에서 클래스별 픽셀 수가 차이나는 것을 보정하고, 균형을 맞춰주는 역할이다. 식3에서 wc(x)는 클래스별로 다른 상수값을 가지는데, 픽셀이 많은 클래스는 낮은값, 픽셀이 적은 클래스는 높은 값이다. 그림 5(b)에서 배경보다 세포가 더 많이 나타난다. 이에 따라 그림 5(d) weight map에서 배경이 세포보다 좀더 밝은 파란색을 띠는데 더 높은 가중치 값을 갖는다는 의미다.
Q : distance는 어떻게 구하는가. L1, L2, ...
또 다른 역할은 같은 클래스의 개체끼리 붙어있을 때, 개체를 구분하는 것이다. 그림 6에서 경계 가중치항만 따로 살펴보자. 오른쪽 그래프는 sqrt(1/e)(x) 함수 그래프다. 이 함수는 입력값(x)이 클수록 가중치는 작아지고, 입력값이 작을수록 가중치는 커지는 특징이 있다.
sqrt(1/e)(x) 함수의 입력값 (D/σ)²은 어떻게 정해지는지 살펴보자. d₁(x)와 d₂(x)를 합한 D는 근접한 두 세포와의 거리합이다. 경계 가중치항에서 (D/σ)²은 항상 0 보다 크거나 같다. 따라서 w0을 제외한 경계 가중치항은 0과 1 사이값을 갖는다. 본 논문에서는 w0을 10으로 정했기 때문에 전체 경계 가중치 항은 0와 10사이의 값을 갖는다.
D와 가중치는 비례하고, σ와 가중치는 반비례한다. 인접한 두 세포와의 거리가 가까운 픽셀은 세포를 구분하는 경계라고 할 수 있고, 이런 픽셀일수록 가중치값이 크다는 것이다. σ는 5 정도로 설정했는데, 이는 근접한 두 세포와의 거리합이 5픽셀보다 가까우면 D/σ이 1보다 커지고, 5픽셀보다 멀면 D/σ이 1보다 작아진다는 것이다. 입력값은 D/σ을 제곱한 (D/σ)²이기 때문에 D가 5픽셀보다 가까우면 입력값이 더 작아지고, 5보다 멀면 더 커진다.
Q : 손실함수 = -에너지 함수 ?
3.3 Data Augmentation
학습할 데이터가 굉장히 적을 때, Data Augmentation에서 중요한 것은 불변성(invariance)와 강인성(robustness)을 보장하는 것이다. 구체적으로는 shift와 회전(rotation)에 대해 불변해야 하고, 왜곡(deformation)과 gray value variation에 강인해야 한다. (확실한 이해를 위해서는 [2]를 참고해야 한다.) 개인적으로는 이미지 데이터를 변형해도 입력 데이터의 분포가 크게 변하지 않는 특성을 강조하는 듯 싶다.
본 논문에서는 3x3 무작위 변위 벡터(random displacement vector)를 이용해 부드럽게 왜곡된 데이터를 얻는다. 변위 값(displacements)은 10(픽셀)표준편차의 가우시안 정규분포에서 샘플링된다. 픽셀별 변위값은 bicubic interpolation(쌍삼차 보간법)으로 계산된다. contracting path 끝에 있는 Drop-out 층들은 추가적인 data augmentation 역할을 수행한다.
- Q1 : random displacement vector란?
- Q2 : 변위값 샘플링?
- Q3 : 쌍삼차 보간법?
- Q4 : Drop out의 Data augmentation 역할?
4. Performance
작성 중
5. Reference
[1] Ciresan, D.C., Gambardella, L.M., Giusti, A., Schmidhuber, J.: Deep neural networks segment neuronal membranes in electron microscopy images. In: NIPS. pp. 2852–2860 (2012)
[2] Long, J., Shelhamer, E., Darrell, T.: Fully convolutional networks for semantic segmentation (2014), arXiv:1411.4038 [cs.CV]
[3] Hyeonwoo Noh, Seunghoon Hong, Bohyung Han. Learning Deconvolution Network for Semantic Segmentation. In 2015.
[4] U-Net 공식 사이트. Uni Freiburg, Computer Vision. site : https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/
[5] Dosovitskiy, A., Springenberg, J.T., Riedmiller, M., Brox, T.: Discriminative unsupervised feature learning with convolutional neural networks. In: NIPS (2014)
