공부/기계학습 및 인공지능

U-Net 간단 리뷰

snn.il 2021. 9. 14. 23:13

U-Net

→ Biomedical 분야에서 이미지 분할(Image Segmentation)을 목적으로 제안된 End-to-End 방식의 Fully-Convolutional Network 기반 모델

  • Biomedical Image Segmentation을 위한 네트워크
    → Biomedical Image 특성 상 전체적인 사진의 부분 부분의 클래스 도출을 위함
  • U-Net은 FCN을 기반으로 확장한 네트워크
    → resolution을 증가, Skip connection 사용
  • 더 적은 데이터를 가지고 더욱 정확한 Segmentation을 수행
    → Data augumentation기법을 사용

[Network Architecture]

→ U-Net은 이미지의 전반적인 컨텍스트 정보를 얻기 위한 네트워크와 정확한 지역화(Localization)를 위한 네트워크가 대칭 형태로 구성

<Contracting Path (Encoder)>

Contracting Path

  • 3x3 convolutions을 두 차례씩 반복 (패딩 없음)
  • 활성화 함수는 ReLU
  • 2x2 max-pooling (stride: 2)
  • Down-sampling 마다 채널의 수를 2배로 늘림

<Expanding Path (Decoder)>

Expanding Path

  • 2x2 convolution (Up-Sampling)
  • 3x3 convolutions을 두 차례씩 반복 (패딩 없음)
  • Up-Conv를 통한 Up-Sampling 마다 채널의 수를 반으로 줄임
  • 활성화 함수는 ReLU
  • Up-Conv 된 특징맵은 Contracting path의 테두리가 Cropped된 특징맵과 concatenation 함
  • 마지막 레이어에 1x1 convolution 연산
  • Convolution 연산에서 패딩을 사용하지 않았기 때문에 최종 출력인 Segmentation map의 크기는 Input Image 크기보다 작음

✔ 이러한 CNN 네트워크의 Feature hierarchy의 결합을 통해 Segmentation이 내제하는 Localization과 Context(Semantic Information) 사이의 트레이드 오프를 해결할 수 있다.

[Training]

<Data Augmentation>

→ Biomedical 분야 특성 상 적은 데이터로 인해 Data Augmentation이 필수

  • Elastic Deformation (이미지에 약간의 변형을 주어 Augmentation 진행)
  • 자연스러운 distortion을 표현하기 위해 각 픽셀에 noise 추가

<Overlap-Tite Input>

→ Fully Convolutional Network 구조의 특성상 입력 이미지의 크기에 제약이 없기에 크기가 큰 이미지의 경우 이미지 전체를 사용하는 대신 overlap-tite 전략을 사용

Overlap-Tite 전략 / 미러링을 이용한 Extrapolation 기법

  • 이미지를 타일형태로 나눠서 입력으로 사용 (파란색과 빨간색 부분)
  • tile에 대한 Segmentation을 얻기 위해서는 이전 입력의 일부분이 포함되어야 함 (Overlap-Tite 전략)
  • 이미지의 경계 부분 픽셀에 대한 세그멘테이션을 위해 0이나 임의의 패딩값을 사용하는 대신 이미지 경계 부분의 미러링을 이용한 Extrapolation 기법을 사용

<Touching cells separation>

  • 각 세포 사이의 경계를 포착할 수 있어야 함
  • 학습 데이터에서 각 픽셀마다 클래스 분포가 다른 점을 고려하여 사전에 Ground-Truth에 대한 Weight map을 구해 학습에 반영
  • → 각 cell의 테두리와 테두리 사이에 반드시 배경이 되도록, 즉 둘 사이에 반드시 배경으로 인식되게 끔 틈을 만들어 학습을 시킴

<Training>

  • 네트워크의 출력 값은 픽셀 단위의 Softmax로 예측
  • Loss Function은 Cross-Entropy 함수 사용
  • → 다만, Touching Cells 분리를 고려하기 위해 Weight map Loss가 포함
  • Momentum : 0.99 로 설정, Gaussian distribution을 통해 학습 파라미터 초기화

[Conclusion]

  • U-Net은 FCNs보다 확장된 개념의 Up-sampling과 Skip Architecture를 적용한 모델을 제안
  • 결과적으로 U-Net의 구조는 아주 적은 양의 학습 데이터만으로 Data Augmentation을 활용하여 여러 Biomedical Image Segmentation 문제에서 우수한 성능을 보임

★) FCN과의 차이점

  1. Upsampling 과정에서 feature channel의 수가 많다는 점
    → contracting path에서의 layer의 정보를 Expansive path의 layer에 skip connection을 적용했기 때문
  2. 각 Convolution의 valid part만 사용
    → Overlap-tile 기법을 사용하여 원활한 segmentation이 가능하도록
    → 사진처럼 비어있는 부분을 mirroring을 통해 다른 변수와의 관계에 기초하여 변수값을 추정

★) SegNet과의 차이점

  1. Encoder의 feature map을 그대로 crop 해오는 Decoder 부분이 SegNet보다 속도가 느리게 만듦
  2. 하지만 이로 인한 성능 향상으로 SegNet보다 좋은 성능을 보임