U-Net : Convolutional Networks for Biomedical Image Segmentation (2015.05.)
Olaf Ronneberger, Philipp Fischer, Thomas Brox
https://arxiv.org/abs/1505.04597
오늘 리뷰할 논문은 의료 이미지 Segmentation에 초점을 맞춰 작성되었지만, 방법론상 많은 강점을 가지고 있어
일반적으로 Segmentation 전 분야에서 현재까지도 널리 활용되는 Architecture인 U-Net : Convolutional Networks for Biomedical Image Segmentation입니다. 논문 리뷰 시작하겠습니다! 😁
📌 Abstract & Introduction
위의 그림은 기존에 일반적으로 활용되던 CNN Architecture를 도식화 한 것입니다. 다양한 변형 CNN들이 있지만 이러한 구조를 띄는 CNN들의 특징은 입력이 최종 말단부의 Fully Connected Layer를 거쳐, 그 출력이 클래스 레이블의 확률 분포로 나온다는 것입니다.
이 방식 자체가 단순 분류 문제에는 적합할 수 있으나, 본 논문에서의 의료 이미지 처리와 같이 특정 상황에서는 label prediction task와 더불어 localization task까지 수행해야 하기 때문에 Image Segmentation Task에 적용하기에는 별도의 구조가 필요합니다. 변형된 세포가 A세포인지, B세포인지도 중요하지만 그 세포가 어디 위치에 있는지 어디까지를 그 세포로 봐야하는지도 중요하니까요!
Cf) Image Segmentation?
Image segmentation란 이미지를 pixel by pixel 별로, 각 pixel이 어느 클래스 카테고리에 속하는지 분류하는 Task 입니다. 객체의 카테고리 끼리만 영역 구분을 하면 Semantic Segmentation, 같은 카테고리 내에서 객체 하나하나 까지 구분하면 Instance Segmentation으로 부릅니다.
[슬라이딩 윈도우 방식]
U-Net 이전에도 Localization 수행을 위해 슬라이딩 윈도우 기반 방식이 설계되었던 적이 있습니다. 전체 이미지를 패치 단위로 분할하고, 패치의 중앙 픽셀의 레이블을 주변부 픽셀들을 통해 예측하고 패치를 Stride에 따라 이동시켜가며 모든 패치의 중앙 픽셀들의 class를 예측하는 과정입니다.
픽셀 단위로 정보를 예측하며, 이미지를 패치 단위로 분할하여 학습에 활용하기 때문에 학습 데이터가 많아진다는 장점이 있습니다. 또한, 패치의 좌표에 따라 위치 정보를 파악할 수 있기 때문에 localization 자체가 이론상 가능합니다.
하지만, 패치 사용의 명확한 한계도 존재합니다. 네트워크가 패치마다 실행되어야 하고, 패치가 독립적으로 계산되어 패치 기준 상대적으로 경계에 있는 정보들은 불안정할 수 있는데, 이를 개선하기 위해 패치 경계를 일부 겹치는 Overlap 전략을 사용할 수 있습니다. 그러나 이 방식의 경우, 중복 계산이 많아져 비효율성이 증가됩니다.
또한, 패치를 크게 하면 전역 정보를 상대적으로 더 많이 담게 되어 문맥 정보를 얻을 수 있지만 세밀한 위치 정보의 손실이 발생하고, 패치를 작게 하면 세밀한 위치 정보나 경계는 잘 파악하지만 문맥 정보의 손실이 생깁니다. 일종의 Trade-Off가 있습니다.
이 문맥 정보와 세부 위치 정보 간 Trade-Off를 극복하고 두 가지를 동시에 고려할수 있는 밸런스 있는 모델을 통해 Segmentation을 수행하고자 Fully Convolutional Network, Skip connection과 Upsampling이 동시에 적용된 바로 본 논문의 architecture, 대칭구조를 띄는 U-Net이 등장하게 되었습니다.
[Overlap tile Strategy & Mirroring]
사실 U-Net은 전체 이미지를 모델에 태울 수 있지만, 고해상도의 의료 이미지를 처리하는 과정에서 메모리 한계상 전체 이미지를 한번에 처리하는 것이 쉽지 않을 수 있습니다. 이때는 위의 방법론에서처럼 U-Net에서도 전체 이미지를 패치로 나누어 모델에 태우는 방법론을 선택할 수 있습니다. 이때 패치 간 일부 겹침, 오버랩을 통해 패치 간 문맥 정보를 보존하고 패치의 경계를 매끄럽게 할 수 있습니다.
처음에 읽을 때는 똑같이 오버랩을 적용하는 방식이기 때문에 슬라이딩 윈도우 방식에서의 오버랩과, U-Net에서의 오버랩의 차이에 대해서 잘 이해가 가지가 않았었는데요. 나중에 공부해보니 그 차이는 다음과 같습니다.
- 슬라이딩 윈도우 방식 : 패치 간 독립적 처리에 주목할 필요성이 있습니다. 패치 분할 시 그냥 이미지를 일정 크기 단위로 나누고, 각 패치 내부에서 자체적으로 똑같이 CNN을 통해 정보를 처리합니다. 하지만 그 대신, 문맥 정보를 보정하기 위해서 여러 패치에 슬라이딩 윈도우 방식으로 스트라이드를 조정해 오버랩을 수행할 수 있는데, 이러한 경우 CNN의 중복 적용 픽셀이 늘어나 계산 비효율성이 생깁니다.
- U-Net의 방식 : 애초에 패치를 분할할 때 오버랩이 된 상태에서 분할이 되어지기 때문에 문맥 정보를 보정할 수 있다는 이점은 똑같지만 연산이 애초에 합쳐진 픽셀에 적용되기 때문에 중복 연산을 최소화 할 수 있습니다. 또한 병렬 처리가 용이하다는 장점이 있습니다.
추가적으로, 패치 단위로 이미지를 처리할때 경계 부분의 가장자리 픽셀은 패치 중앙부 픽셀에 비해 인접 픽셀이 적어 문맥 정보가 부족하기 때문에 보정을 위해 주위에 패딩을 추가하는 방법론을 주로 사용합니다. 본 논문에서는 아래 그림과 같이 기존 이미지와 대칭인 미러링 패딩 방식을 적용해 원래 이미지와 유사한 환경을 생성하도록 돕습니다. 이는 오버랩 패치 분할 방식에서 활용됩니다.
또한 업샘플링을 하는 과정에서도 활용됩니다. 이미지를 원래 해상도로 복원할 때 Upconvolution 연산을 하는 과정에서 가장자리 주위에 패딩을 해야하는데 이때 제로 패딩을 하면 불필요한 노이즈가 가미되고, 그렇다고 Skip Connection을 통해 받는 정보는 다운샘플링 과정에서 이미 위치 정보가 소실된 상태이기 때문에 보정할 때 한계가 있습니다. 따라서 이때 가장자리 픽셀의 패딩 방식으로 미러링을 활용합니다.
[Data augmentation]
본 논문에서는 적은 학습 데이터 보완을 위해 이미지의 다양한 변형을 통한 데이터 증강을 시도하였습니다. 논문에서 언급한 증강 방식은 Shift, Rotation, Gray value, Elastic Deformation 인데요, 그 중에서도 특히 Elastic deformation은 아래 보이시는 그림과 같이 Pixel이 랜덤하게 다른 방향으로 뒤틀리도록 변형하는 방식이라고 합니다.
📌 Network Architecture
[U-Net 전체 구조]
위 그림은 U-Net의 전체 도식입니다. 모델은 U자 형태의 대칭 구조를 띄고 있으며, Contracting Path(수축 경로)와 중간 전환 Bottleneck, 그리고 Expanding Path(확장 경로)로 이루어져 있습니다.
[세부 구조]
전체 모델 도식을 기준으로 왼쪽 부분에 해당하는 이미지의 전반적인 다운 샘플링 과정이 수축경로에 해당합니다. 각 층마다 3*3 Conv 2번을 수행하고 각각 ReLU 활성화 함수를 적용한 후 2*2 Max Pooling을 통해 피쳐맵 크기를 지속적으로 줄여나갑니다. 레이어가 깊어질수록 피쳐맵이 더 큰 Receptive Field를 가지기 때문에 전역 문맥 정보를 다룰 수 있지만 다운 샘플링 과정에서 일부 위치 정보가 손실됩니다.
작아진 피쳐맵에서 위치 정보를 복원하여 localization을 수행하기 위해 구조의 오른쪽에 해당하는 확장 경로를 배치합니다. 층마다 3*3 Conv 2번을 수행하는 것은 동일하지만, 다운샘플링 대신 2*2 Up-convolution을 통해 업샘플링을 수행합니다. U-Net의 코드를 보면, nn 모듈의 ConvTranspose2d 메소드를 활용하여 업샘플링을 구현하는데, 아래의 추가 링크에 U-Net에서의 Upsampling 연산의 원리에 대해 자세히 설명되어 있으니, 더 궁금하신 분들은 참고하시기 바랍니다.
[ pytorch ] ConvTranspose2d 알아보기
- 목차 키워드.- Convolution- CNN- AutoEncoder 들어가며.CNN 이나 AutoEncoder, VAE (Variational Auto Encoder) 등을 구현할 때에 Convolution Layer 을 위해서 Con2d 레이어가 사용됩니다.Convolution Layer 로 구성된 Neural Netwo
westlife0615.tistory.com
업샘플링 과정마다 채널 수가 반으로 줄지만, 그래도 여전히 높은 채널 수를 유지하기 때문에 복원 과정에서 전역 정보를 유지하며 세부 정보의 디테일을 효과적으로 복원할 수 있다고 합니다. 최종 출력값을 얻기 위해 마지막 Layer에서는 1*1 conv를 통해 차원을 조정해줍니다.
각 업샘플링 연산을 수행한 이후 그 출력에 같은 단계 수축 경로의 Output을 Concat하는 방식으로 Skip Connection을 사용합니다. 업샘플링을 하면 단계별로 손실이 생길 수 있는데, 해당 부분을 대칭 위치에 있는 수축경로의 Output으로 보정한다고 생각하면 될 것 같습니다. 예를 들어 첫번째 업샘플링에서는 전역정보가 담긴 피쳐맵을 받아 해상도를 2배로 키우면 전역정보의 손실이 오는데, 수축경로의 마지막 층의 Output이 전역 정보를 많이 담고 있기 때문에 이를 Concat하면, 전역 정보를 유지하면서 업샘플링 연산을 수행할 수 있습니다
📌 Training
GPU 메모리를 최대로 활용하기 위해 배치를 단일 이미지로 축소하였습니다. 또한 Stochastic Gradient Descent를 활용합니다.
논문에서 활용하는 손실함수입니다. 픽셀 별로 계산한다는 점과 Softmax기반으로 log loss를 계산한다는 점은 쉽게 이해할 수 있을 것 같은데, 여기서 주목해야할점은 w(x)인 듯 합니다.
Segmentation Task에서는 위에서 여러번 이야기 하였듯 작은 경계를 제대로 예측하는 것이 중요합니다. 전체 Loss는 w라는 픽셀별 가중치를 각 픽셀별 log 확률 값에 곱하여 합하는 형태인데, w(x) 수식을 보게 되면 픽셀이 경계와의 거리가 작을수록 w(x) 가중치가 커져 전체 Loss에서 해당 픽셀이 차지하는 비중이 커지는 것을 볼 수 있습니다. 이렇듯 논문의 저자들은 모델이 경계 픽셀을 조금 더 잘 학습할 수 있도록 전체 손실 함수를 설계하였습니다.
📌 Experiments
[EM Segmentation Challenge]
모델의 성능 평가를 위해 초파리 유충의 특정 부분을 현미경으로 관찰한 단면 이미지 30장을 Training Set으로 제공받고 Label로는 세포와 막이 분리된 Segmentation map을 받습니다.
위의 결과표를 보면, U-Net은 추가적인 전처리 후처리 없이도 Warping Error에서 타 기존 모델 대비 최고 성능을 달성하였습니다.
Warp Error란 위상 수학과 연결되는 개념으로, 간단하게 얘기하면 픽셀 손실이 똑같더라도 형상이 비슷해야 오류율이 줄어든다는데 정확히 수학적으로 이해하지는 못했지만 Ground Truth와 각 Segmentation 출력을 보면 B가 모양적으로 더 비슷한거 같긴합니다.
[ISBI Cell Tracking Challenge]
세포 이미지 관찰에서도 각각 특정 현미경 이미지 35장, 20장을 활용해서 Training과 Test를 진행하였는데 IoU 평가지표 기준 두 개의 데이터셋 모두에서 U-Net이 가장 좋은 성능을 보였습니다.
📌 Conclusion
본 논문에서는 Upsampling과 Skip Connection을 통해 이미지의 문맥 정보와 세부 위치 정보 두 가지를 동시에 추구하였습니다.또한, 계산 효율성 문제까지 패치 분할을 통해 해결하고자 하였습니다. 마지막으로, Overlap-tile 기법과 Mirroring 등을 통해 이미지 처리 과정에서 발생할 수 있는 문맥 정보 손실, 경계 픽셀의 정보 부족 문제 등에 대해 효과적으로 대응한 연구입니다.
📌 느낀점(논문을 읽으면서..)
저번 프로젝트에서 사람의 얼굴을 Segmentation하는 Task를 수행한 적이 있었습니다. 당시에는 그냥 좋다고 해서 U-Net기반 모델을 정확한 이해없이 사용했던 것 같은데, 그 배경과 원리에 대해서 학습을 하고나니 프로젝트의 결과가 아닌 과정 자체에 대해서 뒤늦게나마 더 잘 이해하게 된 것 같아 좋은 것 같습니다!
📌 References
https://velog.io/@lighthouse97/UNet%EC%9D%98-%EC%9D%B4%ED%95%B4