U-Net

U-Net은 이미지의 모든 픽셀에 클래스를 할당하는 시맨틱 세그멘테이션을 위한 합성곱 신경망(CNN)입니다. 프라이부르크 대학교의 Olaf Ronneberger, Philipp Fischer, Thomas Brox는 2015년 U-Net: Convolutional Networks for Biomedical Image Segmentation에서 이 모델을 소개했습니다. 인코더는 이미지 크기를 줄여 이미지에 무엇이 있는지 파악하고, 디코더는 결과를 원래 해상도로 확장합니다. 스킵 연결은 모든 스케일에서 인코더의 세부 정보를 디코더에 전달하므로 객체 경계가 선명하게 표현됩니다. U-Net은 어노테이션이 달린 현미경 이미지 수십 장만으로도 효과적으로 학습되었고, ISBI 2015 세포 추적 챌린지의 두 부문에서 큰 차이로 우승했으며, 의료 영상 분석의 표준 모델로 자리 잡았습니다. 이후 이 설계는 여러 확산 모델의 노이즈 제거 네트워크로 사용되었습니다.

U-Net의 작동 방식#

네트워크는 대칭적인 두 부분으로 구성되며, U자 형태의 두 경로로 표현됩니다.

• 축소 경로(인코더): ReLU 활성화를 적용한 3×3 합성곱 쌍을 반복하며, 각 쌍 뒤에는 해상도를 절반으로 줄이는 2×2 맥스 풀링이 이어집니다. 채널 수는 각 단계에서 64개에서 1,024개까지 두 배로 늘어나므로 네트워크는 공간 세부 정보를 줄이는 대신 더 넓은 수용 영역과 풍부한 특징을 얻습니다.
• 확장 경로(디코더): 각 단계에서 2×2 업컨볼루션으로 해상도를 두 배 높이고 채널 수를 절반으로 줄입니다. 그 결과를 같은 단계의 인코더 특징 맵과 연결한 뒤 3×3 합성곱 두 개로 결합합니다. 마지막 1×1 합성곱은 각 픽셀의 특징을 클래스 점수로 변환합니다.

원래 네트워크에는 합성곱 레이어 23개와 풀링 단계 4개가 있습니다. 패딩 없는 합성곱을 사용했기 때문에 입력 타일이 572×572이면 388×388 세그멘테이션 맵이 출력되었으며, 대형 이미지는 겹치는 타일 단위로 분할했습니다. 오늘날 대부분의 구현에서는 출력 크기가 입력 크기와 일치하도록 합성곱에 패딩을 적용합니다.

스킵 연결#

풀링은 각 특징의 정확한 위치 정보를 버립니다. 병목 지점에 이르면 네트워크는 어떤 영역에 세포막이 있다는 것은 알지만 세포막이 어떤 픽셀을 지나는지는 알지 못합니다. 스킵 연결은 각 인코더 단계의 고해상도 특징 맵을 U자 경로를 가로질러 디코더로 복사해, 업샘플링 경로의 "무엇" 정보와 인코더의 "어디" 정보를 결합합니다. U-Net은 두 정보를 채널 차원에서 연결하는 반면, ResNet의 잔차 연결은 블록 내부의 특징을 더합니다. 다음 PyTorch 코드는 인코더 단계 하나, 병목 하나, 디코더 단계 하나를 구성합니다.

import torch
from torch import nn

# One encoder level, a bottleneck, and one decoder level of a U-Net
encoder = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1), nn.ReLU())
bottleneck = nn.Sequential(nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding=1), nn.ReLU())
upsample = nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2)
decoder = nn.Sequential(nn.Conv2d(32, 16, 3, padding=1), nn.ReLU())
head = nn.Conv2d(16, 2, kernel_size=1)  # 2 classes, e.g. cell and background

image = torch.randn(1, 3, 128, 128)
skip = encoder(image)  # full-resolution features kept for the decoder
context = bottleneck(skip)  # half resolution, twice the channels, wider context

# Skip connection: concatenate upsampled context with the saved high-resolution features
merged = torch.cat([upsample(context), skip], dim=1)
logits = head(decoder(merged))
print(f"context {tuple(context.shape)} -> merged {tuple(merged.shape)} -> logits {tuple(logits.shape)}")

코드는 context (1, 32, 64, 64) -> merged (1, 32, 128, 128) -> logits (1, 2, 128, 128)을 출력합니다. 업샘플링된 문맥 채널 16개와 스킵 채널 16개가 합쳐져 32개가 되며, 출력에는 입력의 128×128 픽셀 각각에 대한 클래스 점수 두 개가 담깁니다.

데이터가 적은 환경에서의 학습#

전자 현미경으로 촬영한 신경 구조의 원래 학습 데이터셋은 512×512픽셀 이미지 30장으로 구성되었습니다. U-Net은 조직의 굽힘을 모방하는 무작위 탄성 변형을 특히 적극적으로 적용하는 데이터 증강과, 서로 맞닿은 세포 사이의 얇은 배경 간격을 강조하는 가중 교차 엔트로피 손실로 데이터 부족을 보완해 인접 세포를 별개의 영역으로 분리했습니다. 오늘날 의료 세그멘테이션 모델은 마스크의 겹침을 직접 측정하고 픽셀의 몇 퍼센트만 차지하는 작은 구조도 처리할 수 있는 Dice 손실을 사용해 학습하는 경우가 많습니다.

변형 모델과 파생 모델#

• 3D U-Net과 V-Net: 2D 연산을 3D 연산으로 바꾸어 CT 및 MRI 스캔과 같은 볼륨을 분할합니다.
• Attention U-Net과 UNet++: 어텐션 게이트 또는 중첩된 고밀도 연결 스킵 경로를 추가해 디코더에 전달되는 인코더 특징을 정교하게 선별합니다.
• nnU-Net: 새 데이터셋에 맞춰 전처리, 아키텍처, 학습을 자동으로 설정하는 프레임워크입니다. MICCAI 2024 연구에 따르면, 엄격하게 검증할 경우 이 프레임워크에서 잘 조정된 CNN 기반 U-Net이 3D 의료 세그멘테이션에서 최신 Transformer 및 Mamba 아키텍처보다 여전히 뛰어난 성능을 보입니다.
• 확산 U-Net: DDPM과 Stable Diffusion 1.x는 U-Net으로 이미지 또는 잠재 표현의 노이즈를 예측하며, 프롬프트가 생성 과정을 조정할 수 있도록 어텐션 레이어를 추가합니다. Stable Diffusion 1.x의 U-Net에는 약 8억 6천만 개의 파라미터가 있으며, ControlNet은 이 U-Net의 인코더에 연결됩니다. 최신 생성 모델은 U-Net을 확산 Transformer로 대체합니다.

U-Net과 관련 모델 비교#

• 완전 합성곱 네트워크(FCN): 2014년에 나온 FCN은 분류기를 픽셀별 예측 모델로 바꾸고 클래스 점수 맵을 더해 거친 예측과 세밀한 예측을 결합했습니다. U-Net은 모든 단계에서 특징을 연결하는 완전한 디코더를 구축했습니다.
• 오토인코더: 인코더와 디코더 구조를 사용하지만 좁은 병목을 거쳐 입력을 재구성하도록 학습합니다. 스킵 연결을 사용하면 오토인코더가 병목을 우회할 수 있지만, U-Net은 정렬된 픽셀별 레이블을 생성하기 위해 스킵 연결이 필요합니다.
• 인스턴스 세그멘테이션 모델: U-Net은 픽셀마다 클래스 하나를 출력하므로, 학습이나 후처리로 분리하지 않으면 같은 클래스에 속한 두 개의 맞닿은 객체가 동일한 레이블을 공유합니다.

Ultralytics YOLO26 시맨틱 세그멘테이션 모델은 -sem 접미사로 표시되며, 축약된 U-Net 디코더를 사용합니다. 넥은 1/32 스케일의 백본 출력을 두 번 업샘플링하고, 그 과정에서 1/16 및 1/8 스케일의 백본 특징을 연결합니다. 그런 다음 1/8 스케일에서 픽셀을 분류하고 클래스 점수를 8배 업샘플링해 입력 크기에 맞춥니다.

실제 적용 사례#

• 위성 이미지 분석: 토지 피복 지도는 모든 픽셀에 물, 숲, 경작지 또는 건물 레이블을 지정해 산림 벌채와 도시 성장을 추적합니다.
• 산업 검사: 긁힘, 균열, 부식의 픽셀 마스크를 사용하면 바운딩 박스로는 측정할 수 없는 결함의 정확한 크기와 모양을 측정할 수 있습니다.
• 의료 영상: U-Net은 진단과 방사선 치료 계획을 위해 CT 및 MRI 스캔에서 장기, 종양, 병변의 윤곽을 표시하고 현미경 이미지에서 개별 세포를 분리합니다. ISBI 2015 세포 추적 챌린지에서 U-Net은 부분적으로 어노테이션된 HeLa 세포 이미지 20장으로 학습해 평균 IoU 77.5%를 달성했습니다. 두 번째로 우수한 방법의 점수는 46%였습니다.