U-Net(MICCAI 2015) 논문 리뷰

김세현·2023년 1월 28일
1

AI 논문

목록 보기
6/7

U-Net

U-Net은 Biomedical(생의학) 분야에서 Image Segmentation을 목적으로 제안된 End-to-End 방식의 Fully-Convolutional Network 기반 모델입니다.

다르게는 Convolutional Networks for BiomedicalImage Segmentation라고 부르기도 합니다.

End-to-End 방식: 어떤 문제를 해결할 때 필요한 여러 과정을 하나의 신경망을 통해 재배치하는 과정. 데이터를 두 단계로 나누어 각각 네트워크를 구축한 후 학습한 후 그 결과를 합침. 이렇게 하는 이유는 과정을 나누는 것이 성능이 더 좋기 때문.

논문의 큰 특징으로는 2가지가 있습니다.

  • 넓은 범위의 이미지 픽셀로부터 의미정보를 추출하고 그것을 기반으로 각 픽셀마다 객체를 분류하는 아키텍처를 제시.
  • 서로 근접한 객체 경계를 잘 구분하도록 학습하기 위해 Weighted Loss를 제시.

논문의 이름이 U-Net인 이유는 구조가 알파벳 U랑 비슷하게 생겨서 그런 이름이 붙었다고 합니다.


모델 구조

모델의 구조로는 크게 3가지로 나눌 수 있습니다.

  • 넓은 범위의 이미지 픽셀을 보며 의미정보를 추출하는 수축 경로
  • 수축 경로에서 확장 경로로 전환되는 전환 구간
  • 의미정보를 픽셀 위치정보와 결합하여 각 픽셀마다 어떤 객체에 속하는지를 구분하는 확장 경로

모델의 Input은 이미지의 픽셀별 RGB 데이터이고 모델의 Output은 이미지의 각 픽셀별 객체 구분 정보(Class)입니다. Convolution 연산과정에서 Padding을 사용하지 않으므로 모델의 Output Size는 Input Size보다 작아집니다. 예시로 572 × 572 크기의 이미지를 Input으로 사용하면 Output으로 388 × 388 이미지가 생성됩니다.

수축 경로 (Contracting Path)

수축 경로에서 아래와 같은 Downsampling 과정을 반복하여 특징맵(Feature Map)을 생성합니다.

  1. 3×3 Convolution Layer + ReLu + BatchNorm (No Padding, Stride 1)
  2. 3×3 Convolution Layer + ReLu + BatchNorm (No Padding, Stride 1)
  3. 2×2 Max-polling Layer (Stride 2)

수축 경로는 주변 픽셀들을 참조하는 범위를 넓혀가며 이미지로부터 Contextual 정보를 추출하는 역할을 합니다. 3×3 Convolution을 수행할 때 Padding을 하지 않으므로 특징맵(Feature Map)의 크기가 감소합니다.

Downsampling 할때마다 채널의 수를 2배씩 증가시키면서 진행합니다. 즉 처음 Input Channel을 64개로 증가시키는 부분을 제외하면 채널은 1, 64, 128, 256, 512, 1024개로 Downsampling 할때마다 증가합니다.

전환 구간(Bottle Neck)

수축 경로에서 확장 경로로 전환되는 구간입니다.

  1. 3×3 Convolution Layer + ReLu + BatchNorm (No Padding, Stride 1)
  2. 3×3 Convolution Layer + ReLu + BatchNorm (No Padding, Stride 1)
  3. Dropout Layer

마지막에 적용된 Dropout Layer는 모델을 일반화하고 노이즈에 견고하게 만드는 장치입니다.

확장 경로(Expanding Path)

확장 경로에서 아래와 같은 Upsampling 과정을 반복하여 특징맵(Feature Map)을 생성합니다.

  1. 2×2 Deconvolution layer (Stride 2)
  2. 수축 경로에서 동일한 Level의 Feature Map을 추출하고 크기를 맞추기 위하여 자른 후 이전 Layer에서 생성된 Feature Map과 Concatenation합니다.
  3. 3×3 Convolution Layer + ReLu + BatchNorm (No Padding, Stride 1)
  4. 3×3 Convolution Layer + ReLu + BatchNorm (No Padding, Stride 1)

확장경로는 Skip Connection을 통해 수축 경로에서 생성된 Contextual 정보와 위치정보 결합하는 역할을 합니다. 동일한 Level에서 수축경로의 특징맵과 확장경로의 특징맵의 크기가 다른 이유는 여러번의 패딩이 없는 3×3 Convolution Layer를 지나면서 특징맵의 크기가 줄어들기 때문입니다.


학습 방법

논문에서는 다양한 학습 방법등을 통해 모델의 성능을 향상시켰습니다.

  • Overlap-tile strategy : 큰 이미지를 겹치는 부분이 있도록 일정크기로 나누고 모델의 Input으로 활용합니다.
  • Mirroring Extrapolate : 이미지의 경계(Border)부분을 거울이 반사된 것처럼 확장하여 Input으로 활용합니다.
  • Weight Loss : 모델이 객체간 경계를 구분할 수 있도록 Weight Loss를 구성하고 학습합니다.
  • Data Augmentation : 적은 데이터로 모델을 잘 학습할 수 있도록 데이터 증강 방법을 활용합니다.

Overlap-tile strategy

Unet은 GPU memory 자원에 영향을 받지 않고 높은 image resolution 데이터를 사용하기 위해 Overlap-tile 기법을 이용했습니다.

전자 현미경 데이터의 특성상 이미지 사이즈 크기가 매우 커 Patch 단위로 잘라서 Input으로 이용하고, 파란색 영역의 데이터를 입력하여 예측한 노란색 부분의 데이터를 출력합니다.

Mirroring Extrapolate

이미지의 경계부분을 예측할 때에는 Padding을 넣어 활용하는 경우가 일반적입니다. 논문에서는 이미지 경계에 위치한 이미지를 복사하고 좌우 반전을 통해 새 이미지를 생성한 후 원본 이미지의 주변에 붙여 Input으로 사용합니다.

Weight Loss

모델은 위 그림처럼 작은 경계를 분리할 수 있도록 학습되어야 합니다. 따라서 논문에서는 각 픽셀이 경계와 얼마나 가까운지에 따른 Weight-Map을 만들고 학습할 때 경계에 가까운 픽셀의 Loss를 Weight-Map에 비례하게 증가 시킴으로써 경계를 잘 학습하도록 설계하였습니다.

객체간 경계가 전체 픽셀에 차지하는 비중은 매우 작습니다. 따라서 Weight Loss를 이용하지 않을 경우 경계가 잘 학습되지 않아 여러개의 객체가 한개의 객체로 표시 될 가능성이 높습니다.

Data Augmentation

데이터의 양이 적기 때문에 데이터 증강을 통해 모델이 노이즈에 영향을 받지 않도록 학습시킵니다. 데이터 증강 방법으로 Rotation(회전), Shift(이동), Elastic distortion 등이 있습니다.


실험 및 결과

U-Net 모델은 “PhC-U373” 데이터에서 92% IOU Score 를 획득하였으며 2등 모델이 획득한 점수 83% 와 현격한 차이를 보이고 있습니다. U-Net 모델은 “DIC-HeLa” 데이터에서 77.5% IOU Score 를 획득하였으며 2등 모델이 획득한 점수 46% 와 현격한 차이를 보이고 있습니다.

profile
대덕소프트웨어마이스터고 3학년

0개의 댓글