[CV] U-Net

https://arxiv.org/abs/1505.04597

U-Net: Convolutional Networks for Biomedical Image Segmentation

Abstract + Intro

1. 기존 방식의 한계

딥러닝 CV에는 이미지 분류과, 객체 탐지라는 크게 두 가지 Task가 있음

의학 CV 쪽에서는 한 사진 내에서 “어디에” 객체가 있는지와, “어떤” 객체가 있는지.

즉 두 개 모두 높은 성능으로 필요함

• 데이터의 부족: 일반적인 이미지 데이터셋과 달리, 의학 영상은 전문가의 레이블링이 필요하며 양질의 학습 데이터를 대량으로 확보하기가 매우 어려움
• Sliding-Window 방식의 단점: 당시 기존의 CNN 방식은 이미지의 국소 부위를 조금씩 이동하며 중심 픽셀의 클래스를 예측하는 'Sliding-window' 방식을 주로 사용
  • 중복성: 겹치는 패치들을 계속해서 따로 계산해야 하므로 속도가 매우 느림
  • Trade-off: 패치 크기가 작으면 국소적인 특징은 잘 잡지만 전체적인 Context를 놓치고, 패치 크기가 크면 넓은 맥락은 보지만 Localization이 떨어지는 모순

2. 해당 논문의 제안

U-Net은 이러한 한계를 극복하기 위해 FCN 구조를 확장한 독창적인 아키텍처를 제안함

1. 대칭적인 U자형 구조:
   • 수축 경로(Contracting Path - Encoder): 이미지의 Context를 포착
   • 확장 경로(Expanding Path - Decoder): 포착된 특징을 정밀한 해상도로 복원
2. Skip Connections: 수축 경로의 고해상도 특징 맵을 확장 경로의 특징 맵과 직접 결합
3. 이를 통해 인코딩 과정에서 손실된 정교한 위치 정보를 디코더에 전달하여 매우 정확한 경계 추출 가능
4. 강력한 데이터 증강: 적은 데이터로도 학습이 가능하도록 이미지를 Augmentation
5. Overlap-tile 전략: 큰 이미지를 격자 형태로 나누어 처리하면서도 경계 부분의 연속성을 확보하는 기법

3. 성과

• ISBI 챌린지 우승: 2015년 ISBI 세포 추적 챌린지에서 기존의 모든 알고리즘을 큰 차이로 제치고 우승
• 데이터 효율성: 수천 장의 데이터가 아닌, 매우 적은 수의 학습 데이터만으로도 높은 정확도
• 빠른 속도: GPU를 사용하여 512x512 해상도의 이미지를 분할하는 데 1초 미만의 시간이 소요
• 범용성: 특정 도메인에 국한되지 않고 다양한 종류의 의학 영상에서 일관되게 우수한 성능

네트워크 구조

1. 전체 구조

이름 그대로 'U'자 형태의 대칭 구조

왼쪽의 수축 경로(Contracting Path)와 오른쪽의 확장 경로(Expansive Path)로 구성되며

총 23개의 컨볼루션 레이어

2. 수축 경로 (Contracting Path - Encoder)

이미지의 맥락(Context)을 포착하는 전형적인 CNN의 구조

• 수축 경로의 각 단계
  • 3x3 합성곱 → ReLU 활성화
  • 2x2 맥스 풀링 → 특성 맵 크기 절반으로 축소
• 다운샘플링: 2x2 크기, 스트라이드(Stride) 2의 Max Pooling을 사용하여 특징 맵의 크기를 절반으로
• 특징 채널: 다운샘플링을 거칠 때마다 채널의 수는 2배로 증가 (예: 64 → 128 → 256 → 512 → 1024)
• 더 깊은 레이어로 갈수록 더 많은 문맥 정보를 학습할 수 있음

3. 확장 경로 (Expansive Path - Decoder)

포착된 특징을 결합하여 정밀한 위치 정보(Localization)를 복원

• 확장 경로의 각 단계
  • 업샘플링 → 특성 맵 크기 2배로 확대
  • 2x2 업컨볼루션 → 특징 채널 수 절반으로 감소
  • 수축 경로에서 전달된 Skip Connection을 통해 고해상도 정보를 추가로 결합
  • 3x3 합성곱 → ReLU 활성화
• 업샘플링: 2x2 Up-convolution을 사용하여 특징 맵의 크기를 키우고, 채널 수는 절반으로
• 특징 결합: 수축 경로에서 얻은 특징 맵을 확장 경로의 특징 맵과 직접 합침(Skip Connection)
• 연속 컨볼루션: 결합 후 다시 두 개의 3x3 컨볼루션과 ReLU를 적용하여 정보를 정교화

4. 최종 출력 레이어

• 마지막 레이어에서는 1x1 합성곱을 사용해 각 64차원 특징 벡터를 출력 클래스 수에 맞게 변환
• 이를 통해 픽셀 단위의 정밀한 세분화 맵을 생성

5. 주요 기술적 특징

• Skip Connections & Cropping:
  • 수축 경로의 고해상도 정보를 확장 경로에 직접 전달하여 정확한 경계를 찾도록 도움
  • 이때, 패딩이 없는(Unpadded) 컨볼루션을 사용하기 때문에 수축 경로의 특징 맵 크기가 확장 경로보다 큼! 따라서 중앙 부분을 자르는(Cropping) 과정이 반드시 필요
• Final Layer (1x1 Convolution):
  • 네트워크의 마지막 레이어에서는 1x1 컨볼루션
  • 이를 통해 64개의 채널을 사용자가 원하는 클래스 개수에 맞춰 매핑
• Overlap-Tile Strategy:
  • 입력 이미지의 크기에 제한받지 않고, 큰 이미지를 타일 단위로 나누어 처리할 수 있게 설계
  • 타일 경계의 정보 손실을 막기 위해 이미지의 미러링(Mirroring) 기법을 사용
• 패딩 없는 합성곱: 경계 부분의 픽셀 손실을 방지하며, 대형 이미지의 매끄러운 타일링을 가능하게 함

Training

1. 학습 데이터 및 손실 함수

• 입력 데이터: 입력 이미지와 해당 이미지의 정답 분할 맵을 사용
• 손실 함수:
  • 픽셀 단위 소프트맥스
  • 크로스 엔트로피 손실
• 소프트맥스 정의:
  • $a_k(x)$ : 특정 픽셀 $x$ 에서의 $k$ 번째 채널의 활성화 값
  • $p_k(x)$ : $x$ 픽셀이 $k$ 클래스에 속할 확률
• $p_k(x) = \frac{\exp(a_k(x))}{\sum_{k'} \exp(a_{k'}(x))}$
• 크로스 엔트로피 손실 공식:
  • $l(x)$ : 픽셀 $x$ 의 실제 레이블
  • $w(x)$ : 특정 픽셀에 더 높은 중요도를 부여하는 가중치 맵
• $E = - \sum_{x \in \Omega} w(x) \log(p_{l(x)}(x))$
• 최적화: 모멘텀 +SGD
• Batch Size: 높은 해상도의 이미지를 사용하기 때문에 배치 크기를 크게 가져가는 대신, Input Tile의 크기를 크게 설정하여 메모리 효율을 높임 (Batch size 1 기준)

2. 가중치 맵

• 경계 복원을 강화하기 위해 가중치 맵을 사용
• 두 가지 목적:
  1. 클래스 불균형 보정: 특정 클래스가 과소/과대 표현되는 문제 완화
  2. 경계 구분 강화: 접촉하는 객체의 경계 픽셀에 높은 가중치 부여

가중치 맵 정의:

$w(x) = w_c(x) + w_0 \cdot \exp\left(- \frac{(d_1(x) + d_2(x))^2}{2\sigma^2} \right)$

• $w_c(x)$ : 클래스 불균형 보정을 위한 가중치
• $d_1(x)$ : 가장 가까운 객체 경계까지의 거리
• $d_2(x)$ : 두 번째로 가까운 객체 경계까지의 거리
• $w_0$ = 10, 로 설정
• $\sigma \approx 5$

이 방식을 통해 경계 픽셀의 중요도를 높여 경계 복원 성능을 향상시킵니다.

3. 네트워크 초기화

• 깊은 네트워크에서는 가중치 초기화가 매우 중요
• 초기화가 잘못될 경우 일부 경로는 과도한 활성화가 발생하고, 다른 경로는 학습이 비활성화될 수 있음
• U-Net은 He 초기화 방법을 사용해 각 레이어의 활성화 분산이 단위 분산을 유지하도록 설정

He 초기화 공식:

$\text{std} = \sqrt{\frac{2}{N}}$

• $N$ : 한 뉴런에 유입되는 노드 수
• 예: 3x3 합성곱 레이어에서, 이전 레이어의 채널 수가 64일 경우

$N = 3 \times 3 \times 64 = 576$

4. GPU 메모리 최적화

• 패딩 없는 합성곱 사용으로 인해 출력 이미지는 입력 이미지보다 작아짐
• GPU 메모리 사용을 최적화하기 위해 Batch Size 는 1로 설정
• 대신, 모멘텀을 0.99로 높게 설정해 과거의 학습 샘플 정보가 업데이트에 더 큰 영향을 주도록 설계

Data Augmentation

1. 데이터 증강의 필요성

• 생물의학 영상 처리에서는 주석 처리된 데이터가 제한적인 경우가 많음
• U-Net은 다양한 변형을 통해 네트워크가 불변성과 강건성을 학습하도록 설계됨
• 특히, 미세한 조직이나 세포와 같은 현미경 이미지에서는 이동, 회전, 변형 및 명암 변화에 대한 강건성이 중요

2. 탄성 변형

• 탄성 변형은 U-Net의 데이터 증강에서 핵심 기법
• 적은 데이터로도 모델이 강건성을 학습하게 함
• 이는 실제 조직이나 세포에서 나타나는 불규칙한 변형을 효과적으로 반영
• 다양한 형태의 변형을 학습하게 하여 네트워크의 일반화 성능을 높임

구현 방법

1. 3x3 격자 상에서 무작위 변위 벡터를 생성
2. 변위 벡터는 표준편차 10픽셀을 갖는 정규분포에서 샘플링
3. 각 픽셀의 위치는 바이큐빅 보간을 통해 부드럽게 변형

3. 전통적 기법

회전, 평행 이동뿐만 아니라 의학 영상 특유의 노이즈에 대비한 회색조 변화 등을 적용

4. 드롭아웃 활용

• U-Net에서는 드롭아웃 레이어를 통해 암시적 데이터 증강 효과를 추가
• 드롭아웃은 네트워크가 특정 뉴런에 과도하게 의존하는 것을 방지하고, 보다 일반화된 특성을 학습하도록 도움!