비전공생의 UNet++(2018) 논문 리뷰

2021년 나온 SegFormer가 성능이 가장 좋아서 이전에 리뷰한 이후 model 에 대한 논문은 딱히 읽지 않았는데 그래도 유명한 모델이라 한번 읽어보려고 한다. 읽으니 새로운 구조와 정보들이 재밌었다.

 

https://arxiv.org/pdf/1807.10165.pdf

 

Introduction

이 논문 이전까지 SOTA를 이룬 구조는 U-Net과 FCN이라고 한다. 그리고 이러한 Encoder-Decoder 구조의 segmentation모델들은 중요한 공통점을 공유하는데 그것은 skip-connection이다. skip connection은 decoder에서 나온 deep, semantic, coarse-grained feature maps과 encoder에서 나온 shallow, low-level, fine-grained feature map들을 조합해주는 역할을 해준다고 한다. 

위는 UNet의 구조이며 skip connection으로 encoder-decoder의 feature map들이 합쳐지는 것을 알 수 있다. 

 

그러나 natural 한 이미지를 segmentation하는데 UNet은 충분할지 몰라도 medical 분야에서 high level accuracy를 요구하는 segmentation에서는 부족하다고 하며 이는 모델에 대한 불신과 부정확한 병명진단을 할 수 있다고 한다.

 

따라서 저자들은 이러한 문제를 보완하기 위해 UNet++ 라는 구조를 만들었고 nested, dense skip connection이 특징이라고 한다. 이러한 특징들을 통해 더 효과적으로 fine-grained 특징들을 capture 할 수 있었다고 한다.

 

Related Work

FCN은 skip connection을 통해 upsampled 된 feature map들을 sum 해주며, UNet은 이와 다르게 concat을 해준다고 한다. 이전에도 DenseNet 구조에 기인하여 H-DenseNet이라는 구조가 있었지만 이러한 구조는 semantically dissimilar feature maps 을 합치기 때문에 오히려 segmentation 성능이 떨어졌다고 한다.

 

Network Architecture: UNet++

구조는 위 이미지에 잘 나와있으며, (a)가 전체적인 구조이며 (b)에서는 계산하는 방식이 나와있다. 조금 복잡할 수 있는데 (a) 에서 나온 node와 bridge 들을 색깔별로 한번 뜯어보자면, 검은색은 일반적인 UNet구조와 같다. 그리고 X_(i,j)에 대한 연산을 설명해보자면 j부분이 0이라면 그냥 Convolution 연산을 한다. X_(1,0), X_(2,0), X_(3,0) 들을 보면 CNN연산만을 거쳤다. 그리고 X_(0,2)을 본다면 이건 j가 0이 아닌 2인데, 그렇다면 downsampling된 feature에 대해 upsampling된 feature map들과, 그 전 같은 위치해 있는 X_(0,0), X_(0,1)에 대해 concat을 해주고 upsampling, concat 한 feature들에 대해 CNN을 적용한다. 조금 헷갈릴 수 있는데 아래의 수식을 보면 된다.

H는 CNN연산이고 u는 upsampling연산을 나타낸다. 결국 기존 UNet과 다른 점을 요약하자면 단순히 skip connection의 개념만을 사용하는 것이 아니라 여러가지 feature map들을 dense 하게 skip connection을 활용해서 합쳐주는 것이다. 그리고 이 논문에서 deep supervision이란 말이 나오는데, (a)이미지에서 보면 빨간색 선에 해당된다고 한다. 빨간색 선을 통해 1) all segmentation branches are averaged 2) segmentation map is selected from  only one of the segmentation branches 이렇게 표시되어있는데, 무언가 굉장한 뜻을 내포하고 있기 보다는 딘순히 deep 한 feature들을 여러가지 연산 거쳤다 라고 나는 이해하였다.

 

UNet++의 loss는 binarcy cross entropy 와 Dice coefficient를 사용한다. 저자들이 말하는 UNet 과 UNet++의 구분되는 점들은 다음과 같다.

1) having convolution layers on skip pathways which bridges the semantic gap

2) having dense skip connections on skip pathways

3) having deep supervision(shown in red) which enables model pruning and achieve performace to using only one loss layer

따로 해석하지 않고 그냥 원문을 적었는데 어려운 용어들은 아니라 쉽게 이해할 수 있다고 생각한다.

 

Experiment

UNet++모델의 성능을 비교하기 위해 UNet, wide Unet구조와 함께 비교했는데 wide UNet는 기존 UNet구조에 파라미터가 UNet++구조와 같은 모델이다. 이는 단순히 UNet++모델이 파라미터가 UNet보다 많아서 성능이 좋은게 아니라는 것을 보여주기 위해 이렇게 비교군을 구성했다고 한다.

 

해당 수치를 정량적으로 보여주는 표이다.

 

---

재밌었던 논문이다. 한번 코드리뷰를 할 생각이다.