U-Net: Convolutional Networks for Biomedical Image Segmentation

Olaf Ronneberger; Philipp Fischer; Thomas Brox

논문 · arXiv 1505.04597

U-Net: 생의학 영상 분할을 위한 합성곱 신경망 (U-Net: Convolutional Networks for Biomedical Image Segmentation)

U-Net: Convolutional Networks for Biomedical Image Segmentation

Olaf Ronneberger, Philipp Fischer, Thomas Brox

게시일 2015-05-18cs.CV arXiv PDF

XYZNOT로 읽기 (전문 번역과 AI 해설)

세 줄 요약

🎯 매우 적은 수의 주석 데이터셋만으로 고정밀 픽셀 단위 생의학 분할 네트워크를 엔드투엔드로 학습합니다.
🔧 컨텍스트를 위한 수축 경로, 정밀 국소화를 위한 스킵 연결 기반 확장 경로, 패딩 없는 유효 합성곱, 탄성 왜곡 데이터 증강 및 분리 경계 가중치 맵을 결합한 대칭적 U자형 합성곱 아키텍처입니다.
📈 ISBI EM 분할 챌린지에서 슬라이딩 윈도우 방식을 능가했으며(warping error 0.000353), PhC-U373에서 0.9203 IOU, DIC-HeLa에서 0.7756 IOU를 달성하여 ISBI Cell Tracking Challenge 2015에서 우승했습니다.

초록

딥 네트워크를 성공적으로 학습시키려면 수천 개 이상의 주석이 달린 훈련 샘플이 필요하다는 것이 일반적인 통념입니다. 본 논문에서는 데이터 증강을 적극 활용하여 사용 가능한 주석 샘플을 더 효율적으로 사용하는 네트워크 구조 및 훈련 전략을 제시합니다. 이 아키텍처는 컨텍스트를 포착하는 수축 경로(contracting path)와 정밀한 국소화를 가능하게 하는 대칭적 확장 경로(expanding path)로 구성됩니다. 우리는 이러한 네트워크가 극소수의 영상만으로도 엔드투엔드로 학습될 수 있으며, 전자현미경 스택의 신경 구조 분할을 위한 ISBI 챌린지에서 기존의 최고 방법(슬라이딩 윈도우 합성곱 신경망)보다 우수한 성능을 보임을 증명합니다. 투과광 현미경 영상(위상차 및 DIC)으로 훈련된 동일한 네트워크를 사용하여 해당 범주의 ISBI 셀 트래킹 챌린지 2015에서도 큰 격차로 우승을 차지했습니다. 나아가 이 네트워크는 매우 빠릅니다. 최신 GPU에서 512x512 크기 영상의 분할은 1초도 걸리지 않습니다. 전체 구현체(Caffe 기반) 및 훈련된 네트워크는 http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net 에서 확인할 수 있습니다.

영어 원문 초록

There is large consent that successful training of deep networks requires many thousand annotated training samples. In this paper, we present a network and training strategy that relies on the strong use of data augmentation to use the available annotated samples more efficiently. The architecture consists of a contracting path to capture context and a symmetric expanding path that enables precise localization. We show that such a network can be trained end-to-end from very few images and outperforms the prior best method (a sliding-window convolutional network) on the ISBI challenge for segmentation of neuronal structures in electron microscopic stacks. Using the same network trained on transmitted light microscopy images (phase contrast and DIC) we won the ISBI cell tracking challenge 2015 in these categories by a large margin. Moreover, the network is fast. Segmentation of a 512x512 image takes less than a second on a recent GPU. The full implementation (based on Caffe) and the trained networks are available at http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net .

논문 해설

문제 및 동기

생의학 영상 분할(Biomedical image segmentation)은 영상 수준의 분류보다는 고밀도의 픽셀 단위 분류(국소화, localization)를 필요로 합니다. 이 분야에는 두 가지 근본적인 어려움이 존재합니다:

주석이 달린 훈련 데이터가 매우 부족하여, 수천·수만 장이 아닌 수십 장 정도에 불과한 경우가 많습니다.
패치 기반의 슬라이딩 윈도우 합성곱 신경망(예: Cireşan et al.)과 같은 기존 딥러닝 기법은 연산 중복이 심하고 추론 속도가 느리며, 패치 컨텍스트와 국소화 정밀도 간의 본질적인 트레이드오프가 존재했습니다.

저자들은 전역 컨텍스트와 세밀한 국소 디테일을 동시에 포착하고, 임의의 큰 이미지를 매끄럽게 처리할 수 있으며, 극히 제한된 데이터로도 성공적으로 학습할 수 있는 엔드투엔드 완전 합성곱 신경망을 구현하고자 했습니다.

핵심 아이디어

대칭적 U자형 아키텍처(Symmetric U-Shaped Architecture): 수축(다운샘플링) 인코더에 맞춰 컨텍스트를 상위로 전달할 수 있는 수많은 피처 채널을 갖춘 대칭적인 확장(업샘플링) 디코더를 결합하여 완전 합성곱 신경망을 확장했습니다.
스킵 연결(Skip Connections): 수축 경로의 고해상도 피처 맵을 복사하여(크롭 후) 확장 경로의 업샘플링된 피처 맵과 연결(concatenate)함으로써 풀링 과정에서 손실된 정밀한 공간 국소화 세부 정보를 보원합니다.
유효 합성곱을 활용한 오버랩 타일 전략(Overlap-Tile Strategy with Valid Convolutions): 패딩이 없는(unpadded) 유효 합성곱을 사용합니다. 타일링 기법과 미러링을 통해 부족한 경계부 컨텍스트를 외삽함으로써 경계 영역을 매끄럽게 예측하고 대형 영상 처리 시 GPU 메모리 한계를 극복합니다.
탄성 왜곡 증강(Elastic Deformation Augmentation): 성긴 랜덤 그리드를 통해 부드럽고 사실적인 탄성 왜곡을 생성하여 네트워크에 변형 불변성을 학습시킴으로써 필요한 훈련 샘플 수를 대폭 줄입니다.
분리 경계 손실 가중치(Separation Border Loss Weighting): 동일한 클래스의 맞닿아 있는 세포 사이 배경 경계 영역의 오류에 더 큰 페널티를 주는 픽셀 단위 사전 계산 가중치 맵을 적용하여 개별 인스턴스 분리를 용이하게 합니다.

방법론

네트워크 구조

이 네트워크는 2개의 경로로 구성된 23개의 합성곱 레이어를 가집니다:

수축 경로(Contracting Path): 2개의 $3 \times 3$ 패딩 없는 합성곱(각각 ReLU 적용)과 스트라이드 2의 $2 \times 2$ 맥스 풀링으로 구성된 블록을 반복합니다. 풀링 단계를 거칠 때마다 피처 채널 수가 2배로 증가합니다.
확장 경로(Expansive Path): 업샘플링 연산에 이어 피처 채널을 절반으로 줄이는 $2 \times 2$ 합성곱("업 합성곱"), 일치하는 해상도의 크롭된 수축 피처 맵과의 결합(concatenation), 그리고 2개의 $3 \times 3$ 합성곱(ReLU 적용)으로 구성된 블록을 반복합니다.
최종 레이어: $1 \times 1$ 합성곱을 사용하여 64채널의 피처 표현을 원하는 $K$ 개의 클래스 점수로 직접 매핑합니다.

손실 함수 및 경계 가중치

네트워크는 픽셀 단위 소프트맥스 교차 엔트로피로 훈련됩니다: $p_{k} (x) = \frac{exp ( a _{k} ( x ))}{\sum _{k^{'} = 1}^{K} exp ( a _{k^{'}} ( x ))}$ $E = x \in Ω \sum w (x) lo g (p_{ℓ (x)} (x))$ 여기서 $a_{k} (x)$ 는 픽셀 $x$ 에서의 피처 채널 $k$ 의 활성화 값을 나타내며, $ℓ (x)$ 는 실제 클래스 레이블, $w (x)$ 는 공간적 손실 가중치 맵입니다.

클래스 빈도 불균형을 해소하고 맞닿은 객체를 분리하기 위해 $w (x)$ 는 다음과 같이 정의됩니다: $w (x) = w_{c} (x) + w_{0} \cdot exp (- \frac{( d _{1} ( x ) + d _{2} ( x ) ) ^{2}}{2 σ ^{2}})$ 여기서 $w_{c} (x)$ 는 클래스 빈도의 균형을 맞추는 항이고, $d_{1} (x)$ 는 가장 가까운 세포의 경계까지의 거리, $d_{2} (x)$ 는 두 번째로 가까운 세포의 경계까지의 거리이며, 저자들은 $w_{0} = 10$ , $σ \approx 5$ 픽셀로 설정했습니다.

최적화

Caffe 환경에서 배치 크기 1(단일 대형 이미지 타일)과 높은 모멘텀(0.99)의 SGD를 사용하여 훈련했습니다.
가중치는 뉴런당 입력 노드 수가 $N$ 일 때 표준편차가 $2/ N$ 인 가우시안 분포로 초기화했습니다.

실험 결과

ISBI EM 분할 챌린지 (신경 구조):
- warping error $0.000353$ , Rand error $0.0382$ , pixel error $0.0611$ 을 기록하며 최고 순위를 달성했습니다(7가지 회전 버전에 대한 평균값 사용, 추가적인 후처리 없음).
- Cireşan et al.의 기존 최고 성능(warping error $0.000420$ , Rand error $0.0504$ )을 뛰어넘었습니다.
ISBI Cell Tracking Challenge 2015:
- PhC-U373 (위상차): 2위 알고리즘의 $0.83$ 대비 평균 IOU(Intersection over Union) $0.9203$ 을 기록했습니다.
- DIC-HeLa (미분 간섭 대비): 2위 알고리즘의 $0.46$ 대비 평균 IOU $0.7756$ 을 기록했습니다.
효율성: 최신 GPU에서 $512 \times 512$ 이미지 분할에 1초 미만이 소요되며, NVIDIA Titan GPU(6 GB) 기준 전체 훈련에는 약 10시간이 소요됩니다.

한계점 및 향후 과제

패딩 없는 합성곱과 크롭: 유효(unpadded) 합성곱을 사용하기 때문에 출력 타일 크기가 입력 타일보다 필연적으로 작아집니다. 따라서 모든 풀링 단계를 고려해 2로 나누어떨어지는 입력 타일 크기를 신중히 선택해야 하며, 이미지 경계에서 미러링이 필요합니다.
휴리스틱 가중치 기반 인스턴스 분리: 맞닿은 객체를 분리하는 방식이 엔드투엔드 인스턴스 임베딩이나 검출 브랜치 대신 수작업 형태학적 가중치 맵에 의존합니다.
2D 중심 설계: 본 연구는 2D 슬라이스에 집중되어 있어, 3D 입체 생의학 영상으로의 확장은 후속 연구 과제로 남겨졌습니다.

핵심 용어

contracting path (수축 경로): 컨텍스트를 포착하고 고수준 피처 표현을 추출하기 위해 패딩 없는 3x3 합성곱과 맥스 풀링을 반복적으로 적용하는 U-Net의 인코더 부분입니다.
expansive path (확장 경로): 피처 맵을 업샘플링하고 채널 차원을 줄이며 정밀한 공간 분할을 재구성하는 U-Net의 대칭적 디코더 부분입니다.
up-convolution (업 합성곱): 업샘플링 후 적용되어 채널 연결(concatenation) 전에 피처 채널 수를 절반으로 줄여주는 2x2 합성곱입니다.
overlap-tile strategy (오버랩 타일 전략): 임의의 큰 이미지를 중첩된 타일로 분할하고 경계 영역을 미러링하여 패딩 없는 합성곱에 필요한 컨텍스트를 제공하는 추론 및 훈련 방식입니다.
elastic deformations (탄성 왜곡): 생체 조직의 형태적 변형을 시뮬레이션하기 위해 쌍입방 보간법을 통해 성긴 그리드 상에 부드러운 변위 벡터를 생성하는 데이터 증강 기법입니다.
weight map (가중치 맵): 클래스 빈도 균형 조정과 맞닿은 객체 분리를 위한 최근접 두 세포 간 거리에 기반한 지수 페널티를 결합한 공간 가변적 손실 가중치 함수입니다.
warping error (워핑 에러): 끊김 없는 신경세포막 보존과 같은 위상학적 정확성을 평가하기 위해 ISBI EM 챌린지에서 사용된 분할 평가 지표입니다.
fully convolutional network (완전 합성곱 신경망): 완전 연결 레이어 없이 합성곱 및 풀링 연산으로만 구성되어 다양한 입력 크기에 대해 조밀한 예측을 수행할 수 있는 신경망입니다.

자주 묻는 질문

U-Net은 왜 스킵 연결에서 피처 맵을 연결하기 전에 크롭(자르기)하나요?

U-Net이 패딩 없는 유효 합성곱(valid convolution)을 사용하기 때문에, 매 합성곱마다 경계 픽셀만큼 피처 맵 크기가 줄어듭니다. 따라서 크기 차이를 맞추고 정확히 결합(concatenation)하기 위해 수축 경로의 피처 맵을 업샘플링된 피처 맵의 크기에 맞춰 크롭합니다.

U-Net은 추론 시 대형 이미지의 경계 영역을 어떻게 처리하나요?

U-Net은 오버랩 타일 전략을 사용합니다. 타일의 유효한 내부 영역을 온전히 예측하기 위해 주변 컨텍스트 픽셀을 입력으로 함께 제공하며, 이미지 경계 밖의 누락된 데이터는 미러링 반사를 통해 외삽하여 채웁니다.

U-Net은 서로 맞닿아 있는 세포들을 어떻게 분리하나요?

손실 함수에 사전 계산된 가중치 맵을 적용합니다. 가장 가까운 세포 경계와 두 번째로 가까운 세포 경계 사이에 위치한 배경 픽셀에 큰 지수 페널티를 부여함으로써 네트워크가 얇은 분리 경계를 강하게 학습하도록 유도합니다.

훈련 중 배치 크기를 1로 설정한 이유는 무엇인가요?

저자들은 사용 가능한 GPU 메모리(6 GB NVIDIA Titan) 내에서 컨텍스트를 극대화하기 위해 다중 배치 대신 큰 단일 이미지 타일을 선택했습니다. 단일 타일 배치로 인한 그래디언트 불안정성은 0.99의 높은 SGD 모멘텀을 적용하여 보완했습니다.

이 페이지의 해설은 AI가 논문을 바탕으로 작성했으며 오류가 있을 수 있습니다. 자세한 내용은 원문을 확인하세요.

XYZNOT로 읽기 (전문 번역과 AI 해설)