U-Net 개요

U-Net 은 “완전 합성곱 네트워크” (Fully Convolutional Network, FCN)에 기반을 둡니다.
이 구조를 수정 및 확장하여 매우 적은 훈련 이미지로 작동하고 더 정밀한 세분화를 생성합니다. (그림 1 참조).

주요 아이디어는 일반적인 축소 네트워크를 연속적인 레이어로 보완하는 것입니다.
이 레이어에서는 풀링 연산자가 업샘플링 연산자로 대체됩니다.

따라서 이러한 레이어는 출력의 해상도를 높입니다.
로컬라이제이션을 위해 축소 경로의 고해상도 특징이 업샘플링된 출력과 결합됩니다.
이어지는 합성곱 레이어는 이 정보를 바탕으로 더 정밀한 출력을 조합할 수 있습니다.

U-Net 설계 배경 – 합성곱 신경망의 단점

합성곱 신경망은 일반적으로 이미지에 대한 단일 클래스 레이블을 출력하는 분류 작업에 사용됩니다. 그러나 생의학 이미지 처리와 같은 많은 시각적 작업에서는 각 픽셀에 클래스 레이블을 할당하는 로컬라이제이션이 필요합니다.

생의학 작업에서는 수천 개의 훈련 이미지를 얻기 어려우므로, Ciresan 등은 각 픽셀 주위의 국소 영역(패치)을 입력으로 제공하여 픽셀의 클래스 레이블을 예측하는 슬라이딩 윈도우 설정에서 네트워크를 훈련했습니다. 이 네트워크는 ISBI 2012에서 EM 세분화 챌린지를 큰 차이로 우승했습니다.

Ciresan 등의 전략에는 두 가지 단점이 명백합니다.

첫째, 네트워크가 각 패치에 대해 별도로 실행되어야 하므로 매우 느립니다.
중복된 패치로 인해 많은 중복이 발생합니다.

둘째, 로컬라이제이션 (localization) 정확도와 컨텍스트 사용 간에 절충이 필요합니다.
더 큰 패치는 더 많은 최대 풀링 레이어를 필요로 하여 로컬라이제이션 정확도가 감소합니다.
반면, 작은 패치는 네트워크가 적은 컨텍스트만 볼 수 있게 합니다.

최신 접근법들은 여러 레이어의 특징을 고려한 분류기 출력을 제안하여 좋은 로컬라이제이션과 컨텍스트 사용을 동시에 가능하게 합니다.

U-Net 아키텍처 개요

U-Net 의 아키텍처에서 중요한 수정 사항 중 하나는 업샘플링 부분에서 많은 수의 특징 채널을 가지고 있다는 것입니다. 이는 네트워크가 문맥 정보를 고해상도 레이어로 전파할 수 있게 합니다.

그 결과, 확장 경로는 축소 경로와 대칭적이며 U자형 구조를 이룹니다.

네트워크는 완전 연결 레이어를 포함하지 않으며, 각 합성곱의 유효한 부분만 사용합니다.
즉, 세분화 맵에는 입력 이미지에서 전체 문맥이 제공되는 픽셀만 포함됩니다.

이 전략은 겹쳐진 타일 전략을 통해 임의로 큰 이미지를 매끄럽게 세분화할 수 있게 합니다. (그림 2 참조)

그림 2. 겹쳐진 타일 전략

이미지의 경계 영역의 픽셀을 예측하기 위해, 부족한 문맥은 입력 이미지를 반사하여 추정합니다.
이 타일링 전략은 GPU 메모리의 제한 없이 큰 이미지에 네트워크를 적용하는 데 중요합니다.
결과적으로, 이 네트워크는 다양한 생물의학 세분화 문제에 적용될 수 있습니다.

이 논문에서는 EM 스택에서 신경 구조를 세분화하는 데 있어 Ciresan 등의 네트워크를 능가한 결과를 보여줍니다. 또한, ISBI 세포 추적 챌린지 2015의 광학 현미경 이미지에서 세포 세분화 결과도 보여줍니다.
여기서 가장 도전적인 두 개의 2D 전달광 데이터셋에서 큰 격차로 우승했습니다.

다음 섹션에서 U-Net 네트워크 아키텍처에 대해서 좀 더 살펴봅니다.

U-Net 네트워크 아키텍처 (Network Architecture)

네트워크 아키텍처는 그림 1에 나와 있습니다.

왼쪽의 수축 경로와 오른쪽의 확장 경로로 구성됩니다.
수축 경로는 전형적인 합성곱 네트워크 아키텍처를 따릅니다.

두 개의 3×3 합성곱(패딩 없는 합성곱)을 반복적으로 적용합니다.
각 합성곱마다 ReLU 활성화 함수와 2×2 최대 풀링 작업 (다운샘플링을 위한 스트라이드 2)을 따릅니다.
각 다운샘플링 단계에서 피처 채널의 수를 두 배로 늘립니다.

확장 경로의 각 단계는 피처 맵을 업샘플링 합니다.
피처 채널 수를 반으로 줄이는 2×2 합성곱(“업-컨볼루션”)을 따릅니다.
수축 경로에서 대응되는 잘린 피처 맵과 연결되며, 두 개의 3×3 합성곱과 각각 ReLU를 따릅니다.
경계 픽셀의 손실로 인해 크로핑이 필요합니다.

최종 층에서는 각 64-컴포넌트 피처 벡터를 원하는 클래스 수로 매핑하는 1×1 합성곱을 사용합니다.
전체적으로 네트워크는 23개의 합성곱 층을 가지고 있습니다.

출력 세그먼트 맵의 원활한 타일링을 위해 입력 타일 크기를 선택할 때 모든 2×2 최대 풀링 작업이 x 및 y 크기가 짝수인 층에 적용되도록 하는 것이 중요합니다.

U-Net 학습

입력 이미지와 해당 세그먼트 맵을 사용하여 확률적 경사 하강법을 통해 네트워크를 훈련합니다.

패딩이 없는 합성곱 때문에 출력 이미지는 입력 이미지보다 일정한 경계 너비만큼 작아집니다. 오버헤드를 최소화하고 GPU 메모리를 최대한 활용하기 위해 대용량 배치 크기보다 큰 입력 타일을 선호하며, 배치 크기를 단일 이미지로 줄입니다. 이에 따라, 이전에 본 많은 훈련 샘플이 현재 최적화 단계의 업데이트를 결정할 수 있도록 높은 모멘텀(0.99)을 사용합니다.

에너지 함수는 마지막 피처 맵에서 픽셀별 소프트맥스를 사용하여 계산되며, 교차 엔트로피 손실 함수와 결합됩니다. 소프트맥스는 다음과 같이 정의됩니다:

$p_k(\mathbf{x}) = \frac{\exp(a_k(\mathbf{x}))}{\sum_{k'=1}^{K} \exp(a_{k'}(\mathbf{x}))}$

여기서 $a_k(\mathbf{x})$ 는 픽셀 위치 $\mathbf{x} \in \Omega$ 에서 피처 채널 $k$ 의 활성화 값을 나타냅니다. $K$ 는 클래스의 수이며, $p_k(\mathbf{x})$ 는 근사된 최대 함수입니다. 즉, $a_k(\mathbf{x})$ 의 최대 활성화를 가지는 $k$ 에 대해 $p_k(\mathbf{x}) \approx 1$ 이며, 다른 모든 $k$ 에 대해 $p_k(\mathbf{x}) \approx 0$ 입니다.

교차 엔트로피 페널티 부여

교차 엔트로피는 각 위치에서 $p_{\ell(\mathbf{x})}(\mathbf{x})$ 의 1로부터의 편차를 다음과 같이 페널티로 부여합니다:

(1) $E = \sum_{\mathbf{x} \in \Omega} w(\mathbf{x}) \log(p_{\ell(\mathbf{x})}(\mathbf{x}))$

여기서 $\ell : \Omega \rightarrow \{1, \ldots, K\}$ 는 각 픽셀의 실제 라벨이며 $w : \Omega \rightarrow \mathbb{R}$ 는 일부 픽셀에 훈련에서 더 많은 중요성을 부여하기 위해 도입된 가중치 맵입니다.

그림 3. HeLa 셀 DIC (differential interface contrast) 현미경 관찰

각 지상 진리 분할에 대해 가중치 맵을 미리 계산하여 훈련 데이터 세트에서 특정 클래스의 픽셀이 나타나는 빈도를 보상하고, 접촉하는 세포 사이에 도입된 작은 분리 경계를 학습하도록 네트워크를 강제합니다(그림 3c 및 d 참조).

세그멘테이션 계산 – 가중치 맵

분리 경계는 형태학적 연산을 사용하여 계산됩니다. 가중치 맵은 다음과 같이 계산됩니다:

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

여기서 $w_c : \Omega \to \mathbb{R}$ 는 클래스 빈도를 균형 있게 하기 위한 가중치 맵이고, $d_1 : \Omega \to \mathbb{R}$ 는 가장 가까운 세포의 경계까지의 거리, $d_2 : \Omega \to \mathbb{R}$ 는 두 번째로 가까운 세포의 경계까지의 거리입니다. 실험에서는 $w_0 = 10$ 과 $\sigma \approx 5$ 픽셀을 사용했습니다.

많은 합성곱 층과 네트워크를 통과하는 여러 경로가 있는 깊은 네트워크에서는 가중치의 초기화가 매우 중요합니다. 그렇지 않으면 네트워크의 일부는 과도한 활성화를 주고, 다른 부분은 전혀 기여하지 않을 수 있습니다.

이상적으로는 네트워크의 각 피처 맵이 대략 단위 분산을 가지도록 초기 가중치를 조정해야 합니다. 아키텍처(합성곱 및 ReLU 층 대체)의 경우 초기 가중치는 표준 편차가 $\sqrt{2/N}$ 인 가우시안 분포에서 가져오며, 여기서 $N$ 은 한 뉴런의 입력 노드 수를 나타냅니다.

예를 들어, 3×3 합성곱과 이전 층에서 64개의 피처 채널에 대해 $N = 9 \cdot 64 = 576$ 입니다.

데이터 증강 (Data Augmentation)

데이터 증강은 적은 수의 학습 샘플만 사용할 수 있을 때 네트워크에 원하는 불변성과 강인성 특성을 가르치는 데 필수적입니다. 현미경 이미지를 처리할 경우 주로 이동 및 회전 불변성뿐만 아니라 변형과 회색 값 변동에 대한 강인성이 필요합니다. 특히 학습 샘플의 무작위 탄성 변형은 매우 적은 수의 주석이 달린 이미지로 세그멘테이션 네트워크를 학습시키는 핵심 개념으로 보입니다.

여기서는 3×3 격자의 무작위 변위 벡터를 사용하여 부드러운 변형을 생성합니다. 변위는 표준 편차가 10픽셀인 가우시안 분포에서 샘플링됩니다. 픽셀별 변위는 삼차 보간법을 사용하여 계산됩니다. 수축 경로 끝에 있는 드롭아웃 레이어는 추가적인 암시적 데이터 증강을 수행합니다.

U-Net – 실험 결과

표 1. EM 세그멘테이션 챌린지 랭킹 (2015-03-06)

U-net은 입력 데이터의 7개 회전 버전을 평균하여 추가적인 전처리나 후처리 없이 워핑 오류 0.0003529(새로운 최고 점수, 표 1. 참조)와 랜드 오류 0.0382를 달성합니다. 이는 Ciresan et al.의 슬라이딩 윈도우 컨볼루션 네트워크 결과보다 크게 개선된 것입니다.

표 2. 세그멘테이션 결과 - ISBI 세포 추적 챌린지 2015

또한 U-net 을 빛 현미경 이미지의 세포 분할 작업에 적용했습니다. 이 분할 작업은 2014년과 2015년 ISBI 세포 추적 챌린지의 일부입니다.

첫 번째 데이터 세트 “PhC-U373” 에서 92%의 평균 IOU(교차 비율)를 달성했으며, 이는 83%를 기록한 두 번째로 좋은 알고리즘보다 훨씬 우수한 성과입니다. (표 2. 참조).

두 번째 데이터 세트 “DIC-HeLa” 에서 평균 IOU 77.5%를 달성했으며, 이는 46%를 기록한 두 번째로 좋은 알고리즘보다 훨씬 우수한 성과입니다.

그림 4. ISBI 세포 추적 챌린지

U-Net – 결론

U-net 아키텍처는 매우 다양한 생의학적 분할 응용 프로그램에서 매우 우수한 성능을 발휘합니다.

탄성 변형 (elastic deformation)을 이용한 데이터 증강 덕분에 주석이 달린 이미지가 매우 적게 필요합니다.
NVidia Titan GPU(6GB)에서 단 10시간의 훈련 시간만 소요됩니다.
Caffe 기반 구현과 훈련된 네트워크를 제공합니다.

U-net 아키텍처는 더 많은 작업에도 쉽게 적용될 수 있습니다.