[Paper Review] U-Net 논문 정리와 구현
** 틀린 개념/구현 다수
딥러닝 스터디에서 U-Net도 하기로 했었는데 너무너무 게으름 피워서 이번 중간고사 휴회기간에 하게 되었다.. 할게 산더미지만 하나를 해도 꼼꼼히 리뷰하고 구현! 까지 해보자..
U-Net 논문링크
https://arxiv.org/pdf/1505.04597.pdf
1. Introduction
도입부에서 저자는 CNN이 긴 시간동안 존재했음에도, 모델의 크기와 사용가능한 training set의 크기의 한계로 더 나아가지 못했다고 말합니다. 그러면서 CNN은 output으로 이미지에 한 label을 부여하는 classification에 주로 사용되어 왔지만, 의료분야 이미지 처리에서는 이미지의 각 픽셀에 label을 부여하는 것이 선호된다고 말합니다. 그러면 이 논문에서 이미지의 각 픽셀마다 label을 부여하는 방법으로 제시된 sliding window 방식을 인용하며 문제점을 제기합니다.
1. 네트워크(모델)이 각 patch마다 돌아야 하기 때문에 느리다.
2. localization accuracy와 use of context의 trade-off가 존재한다. 즉, 더 큰 patch는 localization accuracy를 낮추고, 작은 patch는 모델이 작은 context만을 보게 한다.
구글링 해서 이해한 내용을 첨언하자면 U-net논문에서 인용한 이전의 논문에서는 sliding window방식을 사용합니다.
하지만 U-net에서는 검증이 끝난 부분을 다시 또 검증하는 sliding window방식이 아닌 patch단위로 하기 때문에 속도가 더 빠르다고 합니다.
따라서 이 논문에서는 위의 Fig1과 같이 더 적은 training images로 더 정확한 segmentations를 만들어낼 수 있는, fully convolutional 한 아키텍처를 제시합니다. 본 논문에서 소개할 모델에서는, 연속적인 레이어들로 contracting network를 구성하고, contracting network에 이어진 upsampling operator존재합니다. 이런 upsampling 레이어들이 contracting 된 해상도로부터 output의 해상도를 증가시킵니다. localize 하기 위해서, contracting path로부터 나온 높은 해상도의 feature이 upsampled 된 output과 합쳐집니다. 이렇게 연속적인 convolution 레이어는 이러한 정보들에 기반해서 더 정확하게 학습을 할 수 있다고 합니다.
논문에서 제시하는 아키텍처의 upsampling 부분에서는 feature의 channel수가 많기 때문에, upsampling 부분의에서 위로 올라가면 커지는 resolution을 가지는 layer에서도 context information을 전파할 수 있습니다. 이렇게 위의 Fig1처럼 expansive path(upsampling part와 같은 말인듯합니다)와 contracting part는 대칭을 이루어 U자의 구조를 만들어냅니다. 또한 U-net은 fully connected layer가 존재하지 않기에 input image에서 full context를 사용할 수 있다고 합니다.
위의 Fig2는 Overlap-tile strategy를 논문에서 설명하면서 삽입한 이미지입니다. 자세히 보면 input이미지와 segmentation이 추출된 결과의 사이즈가 다릅니다. Fig2의 왼쪽 사진의 희미한 흰색선 내부 부분만이 segmentation이 진행되는데, 흰색선 외부의 부분은 잘 살펴보면 흰색선을 기준으로 대칭적으로 내부와 똑같은 모습을 하고 있습니다.
이는 contracting part에서 패딩이 없어서 이미지 외곽부분이 잘리는 것을 해결하기 위해서 도입한 mirroring이라고 합니다. 이는 missing context를 채워주기 위한 방법입니다. (참고: worb1605님 네이버블로그 )
논문에서는 가진 데이터의 수가 적었기때문에, elastic deformation을 training set에 도입해서 data augmentation을 진행해 주었다고 합니다. 또
U-net의 목적은 biomedical분야에서 세포들을 segmentation하기 위함이었는데, 이렇게 cell segmentation에서 같은 class를 가진 세포가 맞닿아있을 때, 분류하는 어려움이 있었고, 이를 해결하기 위해 weighted loss를 사용했다고 합니다.
2. Network Architecture
U-net의 전체적인 구조는 위의 Fig1에서 확인할 수 있습니다. contracting path와 expansive path로 나뉘어지고, contracting path는 각 단계에서 3x3크기의 convolution 연산을 padding 없이 두 번 진행하고 ReLU를 통과합니다. 또한 밑의 단계로 내려가면서(down sampling) 2x2 크기, 2 stride의 max pooling 연산을 진행하고 채널을 2배로 늘려주었습니다.
expansive path에서는 2x2크기의 up-convolution 연산을 진행하는데, 이것은 윗 단계로 올라갈수록 채널수를 반으로 줄입니다. 그리고 각 단계에서는 3x3크기의 convolution연산을 진행하고 ReLU를 통과합니다. expansive path의 맨 위단계의 마지막에서는 1x1 convolution연산을 진행하는데, 이것은 마지막에 64 채널에서 원하는 class개수로 채널수를 조정해 주기 위함입니다. 그래서 총 23개의 convolution layer가 존재합니다.
주목해야 할 점은 Fig1에서 회색화살표입니다. contracting path의 각 단계에서는 convolution연산이 끝나고 나면 feature map을 crop 해서, expansive path의 같은 단계에 붙여줍니다. 이것은 각 단계의 convolution연산에서 border pixels의 손실 때문에 필수적으로 진행합니다.
3. Training
U-net은 학습시에 input으로 image와 그 이미지의 segmentation map을 입력해 줍니다. 또한 GPU사용을 극대화해 주기 위해서, Overlap-tile strategy을 적용하면서 큰 patch사이즈를 적용한다고 합니다. 위의 Fig2에서 확인할 수 있듯이, Overlap-tile strategy에서는 파란 크기의 patch(batch)를 입력하면, 모델에 padding이 없기 때문에 노란색 크기의 결괏값이 도출됩니다. 이렇게 U-net에서는 단일 이미지를 여러 patch로 나누어서 입력시킨다고 합니다. 또한 0.99라는 큰 모멘텀을 설정하여 이전에 학습하였던 이미지들이 현재 optimization을 하는데 영향을 주게 했다고 합니다.
또한 energy function이란것이 각 픽셀의 최종 softmax결과에 대해서, cross entropy loss와 함께 계산된다고 합니다.
**잠시 보류**
3.1 Data Augmentation
논문에서는 training 데이터셋이 적다면, 모델이 강건하게 학습하기 위해서는 Data augmentation이 필수라고 말합니다. 이러한 Data augmentation에서 가장 중요한것은 random elastic deformation이라고 합니다. 저자들은 10 픽셀에서의 표준편차를 따르는 가우시안 분포로 displacement를 샘플링하고, 그것을 이용해서 이미지를 변형했습니다.
4. Experiments
U-net을 세가지의 다른 segmentation task에서 평가를 했다고 합니다. 다만 저는 한 가지만 소개하도록 하겠습니다.
저자들은 EM segmentation challenge에서 성능을 평가했다고 하는데요, 512x512 크기의 30개의 전자현미경 사진을 학습했습니다. 각 이미지들은 cell과 membrane 부분의 segmentation ground truth가 주어졌습니다. 이렇게 학습을 진행 후, warping error, rand error, pixel error의 세 가지를 측정해서 성능을 평가했습니다. 그리고 U-net은 EM segmentation challenge은 논문의 초반에서 언급한 sliding window방식보다 더 좋은 성능을 보였다고 합니다.
**구현추가 예정