U-Net 구현

탐지 사례의 U-Net은 RD 맵을 어떤 모양으로 바꾸어 셀마다의 확률을 내는가?

U-Net은 잘 알려진 segmentation 구조이지만(18장), 결과를 해석하려면 이 사례의 입력 크기, 채널 수, 패딩, skip connection, sigmoid 출력을 함께 알아야 한다.

구현은 2채널 RD 맵을 입력으로 받는다. 인코더 4단계에서 채널을 32 → 64 → 128 → 256으로 늘리며 공간 해상도를 절반씩 줄인다. 병목(bottleneck)은 512채널 표현에 Dropout2d를 더해 가장 깊은 추상화를 만든다. 디코더 4단계는 up-sampling으로 해상도를 되돌리며, 인코더의 특징을 skip connection으로 같은 해상도 단계에 가져온다. 마지막 1 × 1 합성곱과 sigmoid가 RD 맵 셀마다 양성 확률을 낸다. 전체 모양은 입력과 같은 64 × 200 격자에서 셀마다 0~1의 값을 내는 dense prediction이다. 학습 파라미터는 모두 약 7.76 M개이다.

U-Net 구조 — 블록 옆의 수는 C × H × W 텐서 모양이고, H는 속도 셀, W는 거리 셀이다. 입력 폭 200은 reflect padding으로 208로 늘렸다가 출력에서 200으로 자른다.

출력은 클래스 하나가 아니라 RD 맵과 같은 크기의 확률 맵이므로, 이 맵을 탐지 결과로 바꾸는 임계값을 따로 정해야 한다. 같은 모델이라도 임계값을 어떻게 정하느냐에 따라 운용 결과가 크게 달라진다.

텐서 모양의 단계별 추적

입력 RD 맵은 인코더, 병목, 디코더를 지나며 아래 표처럼 모양이 바뀐다. 구현은 폭 200을 reflect padding으로 208로 늘린 뒤 최종 출력을 폭 200으로 자른다. 공간 해상도는 절반씩 줄었다가 두 배씩 늘고, 채널 수는 그 반대로 움직인다. 이 대칭 구조 덕분에 신경망은 깊은 층에서 큰 그림(어떤 무늬가 있는가)을, 얕은 층에서 정확한 위치(어디에 있는가)를 함께 배운다.

단계	출력 모양(C × H × W)	역할
입력 RD 맵	2 × 64 × 200	log-magnitude + 위상 두 채널
Enc 1(Conv-BN-ReLU ×2)	32 × 64 × 208	저수준 특징 추출. 이 단계에는 pool이 없어 해상도가 그대로이다
Enc 2(Pool + Conv ×2)	64 × 32 × 104	해상도 1/2, 채널 두 배
Enc 3(Pool + Conv ×2)	128 × 16 × 52	해상도 1/4
Enc 4(Pool + Conv ×2)	256 × 8 × 26	해상도 1/8, 가장 넓은 공간 문맥
병목(Pool + Conv ×2 + Dropout2d)	512 × 4 × 13	가장 깊은 추상화 단계
Dec 4(Up + Concat Enc 4 + Conv ×2)	256 × 8 × 26	해상도 복원 시작, skip으로 위치 정보 결합
Dec 3(Up + Concat Enc 3 + Conv ×2)	128 × 16 × 52	해상도 복원 진행
Dec 2(Up + Concat Enc 2 + Conv ×2)	64 × 32 × 104	해상도 1/2 복원
Dec 1(Up + Concat Enc 1 + Conv ×2)	32 × 64 × 208	패딩한 해상도에 도달
최종 1 × 1 Conv + Sigmoid + Crop	1 × 64 × 200	셀마다 양성 확률

skip connection의 역할

인코더의 down-sampling은 위치 정보를 흐린다. 예를 들어 Enc 4의 픽셀 하나는 입력 RD 맵의 8 × 8 셀을 압축한 결과라서, 이 영역에 표적이 있다는 것까지는 알아도 그 안의 어느 셀이 표적인지는 흐릿하다. 병목에서는 16 × 16 셀 수준으로 더 흐려진다.

skip connection은 같은 해상도의 인코더 특징을 디코더에 바로 넘겨 이 흐림을 메운다. 깊은 표현(무엇이 있는가, semantic)과 얕은 표현(어디에 있는가, spatial)을 같은 단계에서 합쳐, 최종 출력이 RD 맵 해상도 그대로 셀 단위의 위치를 짚게 한다. 그래서 U-Net은 RD 맵 전체의 무늬를 보면서도 셀 하나 단위로 표적을 찍는다.

손실 함수: BCE, focal, Dice

이 사례는 양성 셀이 전체의 0.5 % 이하일 만큼 클래스 불균형이 극단적이다. 이때 보통의 BCE만 쓰면 모든 셀을 0으로 예측해도 손실이 작아져 all-zero 예측이 자주 일어난다(18장). 아래 세 손실은 이 압력을 서로 다르게 푼다.

BCE: 모든 셀의 cross-entropy 평균이다. 음성 셀이 손실의 거의 전부를 차지하므로, 불균형한 상태에서는 양성 셀의 신호가 묻힌다.
focal loss: $-\alpha(1-\hat{y})^\gamma \log \hat{y}$ 꼴(양성 샘플, 이 사례는 $\alpha = 0.75$ , $\gamma = 2$ )로, 이미 잘 분류한 샘플의 기여를 $(1-\hat{y})^\gamma$ 로 줄인다. 어려운 샘플이 받는 기울기가 상대적으로 커져, 양성 셀처럼 드문 샘플도 학습에 반영된다.
Dice loss: $1 - \frac{2|P \cap G|}{|P| + |G|}$ 로 예측 마스크와 정답 마스크의 겹침을 직접 최적화한다. 양성 영역의 크기보다 모양과 위치가 얼마나 겹치는지를 보므로, 양성이 작을 때 안정적이다.

불균형 탐지 문제에서 BCE, focal, Dice 손실 비교 — 왼쪽과 가운데는 양성(y = 1)과 음성(y = 0) 셀 하나의 BCE와 focal 손실(γ = 2, α = 0.75)이다. 오른쪽은 양성 1 %, 음성 99 %인 맵에서 전체 손실을 음성과 양성이 나누는 비율이다.

손실 함수를 고르면 어떤 샘플이 학습에 더 크게 반영되는지가 정해진다. 이 사례에서는 focal과 Dice를 조합해 양성 영역의 학습 압력을 유지한다.

결과 표를 만든 학습 설정은 다음과 같다. optimizer는 Adam, 학습률은 $3\times10^{-4}$ 이고, scheduler는 검증 손실 기준의 ReduceLROnPlateau(patience = 5, factor = 0.5)이다. batch는 32, epoch는 30이다. 정규화로는 병목의 Dropout2d(0.3)를 쓰고, Adam에는 따로 weight decay를 주지 않는다. 학습 스크립트의 기본값은 batch 64, 학습률 $10^{-3}$ 이므로 같은 결과를 재현하려면 명령행에서 지정한다. 학습률은 너무 크면 발산하고 너무 작으면 학습이 멈춘다.