Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation 리뷰

Cao, Hu, Yueyue Wang, Joy Chen, Dongsheng Jiang, Xiaopeng Zhang, Qi Tian, and Manning Wang. "Swin-unet: Unet-like pure transformer for medical image segmentation." In European conference on computer vision, pp. 205-218. Cham: Springer Nature Switzerland, 2022. 

Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation

 

Abstract

Medical image analysis에선 CNNs과 더불어 skip-connection이 더해진 U-net형태의 모델이 대표적으로 자리잡고 있다. 하지만 CNN은 지역적 정보를 추출하는 conv layer때문에 전역 정보를 학습하기 어렵다. 우리는 Swin Transformer구조의 encoder-decoder 형태를 가진 Swin-Unet을 제안한다. 입출력을 4배로 다운/업샘플링하는 이 구조의 모델에서 fully-convolution, hybrid모델을 압도하는 성능을 보였다.

 

우리는 Swin Transformer의 성공에서 영감을 받아 가장 처음으로 transformer 기반의 u-net 구조를 설계하였다. 입력이미지는 중복없는 이미지 패치 토큰으로 분할된 후 인코더로 입력된다. 추출된 feature는 디코더의 patch expanding layer를 통해 업샘플링 되고 segmentation prediction을 위해 encoder와 skip-connection된다.

 

Architecture overview

 

Swin-Unet은 인코더, 보틀넥, 디코더, skip-connection으로 구성되어 있다. 입력 이미지는 4x4 크기의 패치로 나누어지며 Linear embedding을 통해 임베딩 된다. 계층적인 Swin Transformer block을 통과하게 되며 Patch merging에서 down-sampling되면서 차원을 증가시킨다. 디코더에서는 skip connection을 통해 인코더의 feature와 연결되어 spatial information loss를 줄인다.

Swin transformer block은 위와 같이 생겼으며 window-MSA, shifted window-MSA 메커니즘 등이 있는데 이는 Swin transformer 논문을 참조하길 바람, 특별한 방법론이 사용된게 아닌 실험결과에 대한 논문이라 특별한 내용은 더 없는듯

 

discussion

- input size 224, 384에 대한 실험 내용이 있는데 부분적으로는 224가 우세한 결과도 있지만 전체적으론느 384가 성능이 좋은듯 하다. 또한 모델 크기도 tiny, base를 비교한 내용이 있는데, 이 역시 부분적으로 tiny가 우세한 면이 있지만 전체적으로는 base크기의 모델이 우세하다.