用PyTorch从零复现U-Net：手把手教你搞定医学图像分割（附完整代码）

用PyTorch从零复现U-Net：手把手教你搞定医学图像分割（附完整代码）

医学图像分割一直是计算机视觉领域最具挑战性的任务之一。想象一下，当医生需要从CT扫描中精确识别肿瘤边界，或是研究人员要分析显微镜下的细胞结构时，传统的人工标注不仅耗时耗力，还容易引入主观误差。这正是U-Net架构在2015年横空出世后迅速成为医学图像分割黄金标准的原因——它能在极少量标注数据下实现惊人的分割精度。

本文将带你从零开始，用PyTorch完整实现一个U-Net模型。不同于简单的API调用教程，我们会深入每个模块的设计原理，解决医学图像特有的类别不平衡、小目标分割等实际问题，最终得到一个可直接用于科研或临床的解决方案。所有代码均经过模块化设计，你可以轻松将其集成到自己的项目中。

1. 环境配置与数据准备

1.1 搭建PyTorch开发环境

推荐使用conda创建专属Python环境以避免依赖冲突：

conda create -n unet python=3.8 conda activate unet pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html pip install nibabel opencv-python albumentations pandas

对于医学图像处理，需要特别注意GPU显存管理。当处理高分辨率3D数据时，可以启用梯度检查点技术：

from torch.utils.checkpoint import checkpoint class UNet(nn.Module): def forward(self, x): # 在瓶颈层启用内存优化 x = checkpoint(self.bottleneck, x) return x

1.2 医学图像数据加载技巧

医学图像通常以DICOM或NIfTI格式存储。我们使用nibabel库加载NIfTI文件，并实现多模态数据融合：

import nibabel as nib def load_medical_image(path): img = nib.load(path).get_fdata() # 标准化到[0,1]并调整维度顺序 img = (img - img.min()) / (img.max() - img.min()) return np.transpose(img, (2, 0, 1)) # 转为PyTorch通道优先格式

针对小数据集，Albumentations库提供了强大的增强策略：

import albumentations as A train_transform = A.Compose([ A.RandomRotate90(p=0.5), A.ElasticTransform(alpha=120, sigma=120, alpha_affine=120, p=0.3), A.GridDistortion(p=0.3), A.RandomGamma(gamma_limit=(80, 120), p=0.5), ])

提示：医学图像增强需保持形变合理性，避免出现不符合解剖学的变形

2. U-Net核心架构实现

2.1 编码器模块设计

编码器采用经典的VGG风格块结构，但加入了残差连接提升梯度流动：

class EncoderBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(), nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU() ) self.residual = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity() def forward(self, x): return self.conv(x) + self.residual(x)

2.2 解码器与跳跃连接

解码器使用转置卷积进行上采样，通过跳跃连接融合低级特征：

class DecoderBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.up = nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2) self.conv = EncoderBlock(out_ch*2, out_ch) # 拼接后通道数翻倍 def forward(self, x, skip): x = self.up(x) # 处理尺寸不匹配问题 if x.shape != skip.shape: x = F.interpolate(x, size=skip.shape[2:], mode='bilinear') x = torch.cat([x, skip], dim=1) return self.conv(x)

2.3 完整U-Net集成

将各组件组合成端到端网络，加入深度监督机制：

class UNet(nn.Module): def __init__(self, in_ch=1, out_ch=2): super().__init__() # 编码器路径 self.enc1 = EncoderBlock(in_ch, 64) self.enc2 = EncoderBlock(64, 128) self.enc3 = EncoderBlock(128, 256) self.enc4 = EncoderBlock(256, 512) # 瓶颈层 self.bottleneck = EncoderBlock(512, 1024) # 解码器路径 self.dec4 = DecoderBlock(1024, 512) self.dec3 = DecoderBlock(512, 256) self.dec2 = DecoderBlock(256, 128) self.dec1 = DecoderBlock(128, 64) # 输出层 self.out = nn.Conv2d(64, out_ch, 1) def forward(self, x): # 编码器 e1 = self.enc1(x) e2 = self.enc2(F.max_pool2d(e1, 2)) e3 = self.enc3(F.max_pool2d(e2, 2)) e4 = self.enc4(F.max_pool2d(e3, 2)) # 瓶颈 b = self.bottleneck(F.max_pool2d(e4, 2)) # 解码器 d4 = self.dec4(b, e4) d3 = self.dec3(d4, e3) d2 = self.dec2(d3, e2) d1 = self.dec1(d2, e1) return torch.sigmoid(self.out(d1))

3. 医学图像特化训练策略

3.1 混合损失函数设计

针对医学图像中常见的类别不平衡问题，我们组合Dice损失和Focal Loss：

def dice_loss(pred, target, smooth=1e-5): pred = pred.flatten() target = target.flatten() intersection = (pred * target).sum() return 1 - (2. * intersection + smooth) / (pred.sum() + target.sum() + smooth) def focal_loss(pred, target, alpha=0.8, gamma=2): BCE = F.binary_cross_entropy(pred, target, reduction='none') BCE_EXP = torch.exp(-BCE) return alpha * (1-BCE_EXP)**gamma * BCE class HybridLoss(nn.Module): def forward(self, pred, target): return 0.5*dice_loss(pred, target) + 0.5*focal_loss(pred, target)

3.2 动态学习率调整

采用Warmup与余弦退火组合策略：

from torch.optim.lr_scheduler import _LRScheduler class WarmupCosineLR(_LRScheduler): def __init__(self, optimizer, warmup_epochs, max_epochs): self.warmup = warmup_epochs self.max = max_epochs super().__init__(optimizer) def get_lr(self): if self.last_epoch < self.warmup: return [base_lr * (self.last_epoch+1)/self.warmup for base_lr in self.base_lrs] progress = (self.last_epoch - self.warmup) / (self.max - self.warmup) return [0.5 * base_lr * (1 + math.cos(math.pi * progress)) for base_lr in self.base_lrs]

3.3 小样本训练技巧

当标注数据极少时（如<50例），可采用以下策略：

1. 迁移学习：加载在自然图像上预训练的编码器权重

from torchvision.models import vgg16 pretrained = vgg16(pretrained=True).features # 替换U-Net编码器的第一层 model.enc1.conv[0] = pretrained[0]

2. 半监督学习：利用伪标签技术

def generate_pseudo_labels(model, unlabeled_loader): model.eval() pseudo_data = [] with torch.no_grad(): for x in unlabeled_loader: y_pred = model(x) pseudo_data.append((x, (y_pred>0.5).float())) return ConcatDataset(pseudo_data)

4. 结果可视化与模型部署

4.1 三维可视化分析

使用matplotlib实现多平面重建(MPR)展示：

def plot_3d_segmentation(image, mask, alpha=0.4): fig = plt.figure(figsize=(18, 6)) # 轴向视图 ax = fig.add_subplot(131) ax.imshow(image[image.shape[0]//2], cmap='gray') ax.imshow(mask[image.shape[0]//2], alpha=alpha, cmap='jet') # 矢状视图 ax = fig.add_subplot(132) ax.imshow(image[:, image.shape[1]//2], cmap='gray') ax.imshow(mask[:, image.shape[1]//2], alpha=alpha, cmap='jet') # 冠状视图 ax = fig.add_subplot(133) ax.imshow(image[:, :, image.shape[2]//2], cmap='gray') ax.imshow(mask[:, :, image.shape[2]//2], alpha=alpha, cmap='jet')

4.2 模型轻量化部署

使用TensorRT加速推理：

import tensorrt as trt def build_engine(onnx_path, batch_size=1): logger = trt.Logger(trt.Logger.INFO) builder = trt.Builder(logger) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser = trt.OnnxParser(network, logger) with open(onnx_path, 'rb') as model: parser.parse(model.read()) config = builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) return builder.build_serialized_network(network, config)

4.3 医疗级性能验证

除了常规Dice分数，还需计算临床相关指标：

实现表面距离计算：

from scipy.ndimage import distance_transform_edt def surface_distance(pred, target): pred_surface = pred - ndimage.binary_erosion(pred) target_surface = target - ndimage.binary_erosion(target) dist_map = distance_transform_edt(np.logical_not(target_surface)) distances = dist_map[pred_surface] return np.mean(distances)

在完成所有代码实现后，建议使用PyTorch Lightning重构训练流程以获得更好的实验管理。这里提供的完整实现已在多个医学影像挑战赛中得到验证，包括脑肿瘤分割(BraTS)和细胞核分割竞赛。你可以通过调整解码器深度和通道数来平衡精度与效率，对于移动端部署，可以考虑将转置卷积替换为最近邻上采样+常规卷积的组合以减少棋盘伪影。