U-Net 医学图像分割实战：PyTorch 实现与 ISBI 数据集 Dice 系数 0.92 复现

U-Net 医学图像分割实战：PyTorch 实现与 ISBI 数据集 Dice 系数 0.92 复现

在医学影像分析领域，图像分割是疾病诊断和治疗规划的关键步骤。U-Net 凭借其独特的对称编码器-解码器结构和跳跃连接机制，已成为细胞分割、肿瘤检测等任务的黄金标准。本文将带您从零实现一个完整的 U-Net 训练流程，通过 PyTorch 框架在 ISBI 细胞分割数据集上复现 Dice 系数 0.92 的论文结果。

1. 环境配置与数据准备

1.1 安装依赖库

确保已安装最新版 PyTorch 和必要的数据处理工具：

pip install torch torchvision pip install opencv-python scikit-image tqdm

1.2 ISBI 数据集处理

ISBI 挑战赛提供 30 张 512x512 的电子显微镜图像及其标注。我们需要将其转换为 PyTorch 可处理的格式：

import os from skimage import io import numpy as np def load_isbi_data(data_dir): images = [] masks = [] for filename in sorted(os.listdir(data_dir)): if 'image' in filename: img = io.imread(os.path.join(data_dir, filename)) images.append(img) elif 'label' in filename: mask = io.imread(os.path.join(data_dir, filename)) masks.append(mask) return np.array(images), np.array(masks)

提示：建议将原始图像归一化到 [0,1] 范围，并添加通道维度 (H,W) -> (1,H,W)

1.3 数据增强策略

为提升模型泛化能力，采用论文中的弹性形变增强：

from scipy.ndimage import map_coordinates from scipy.ndimage.filters import gaussian_filter def elastic_transform(image, alpha=1000, sigma=30): random_state = np.random.RandomState(None) shape = image.shape dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant") * alpha dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant") * alpha x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing='ij') indices = np.reshape(x+dx, (-1, 1)), np.reshape(y+dy, (-1, 1)) return map_coordinates(image, indices, order=1).reshape(shape)

2. U-Net 模型架构实现

2.1 基础模块设计

构建编码器的下采样模块和对应的解码器上采样模块：

import torch.nn as nn class DoubleConv(nn.Module): """(conv => BN => ReLU) * 2""" 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(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True) ) def forward(self, x): return self.conv(x) class Down(nn.Module): """下采样：MaxPool => DoubleConv""" def __init__(self, in_ch, out_ch): super().__init__() self.mpconv = nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_ch, out_ch) ) def forward(self, x): return self.mpconv(x)

2.2 跳跃连接与上采样

实现解码器的特征融合模块：

class Up(nn.Module): """上采样：转置卷积 => 特征拼接 => DoubleConv""" def __init__(self, in_ch, out_ch, bilinear=True): super().__init__() if bilinear: self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) else: self.up = nn.ConvTranspose2d(in_ch//2, in_ch//2, 2, stride=2) self.conv = DoubleConv(in_ch, out_ch) def forward(self, x1, x2): x1 = self.up(x1) diffY = x2.size()[2] - x1.size()[2] diffX = x2.size()[3] - x1.size()[3] x1 = F.pad(x1, [diffX//2, diffX - diffX//2, diffY//2, diffY - diffY//2]) x = torch.cat([x2, x1], dim=1) return self.conv(x)

2.3 完整 U-Net 架构

整合各模块构建对称网络结构：

class UNet(nn.Module): def __init__(self, n_channels=1, n_classes=1): super().__init__() self.inc = DoubleConv(n_channels, 64) self.down1 = Down(64, 128) self.down2 = Down(128, 256) self.down3 = Down(256, 512) self.down4 = Down(512, 512) self.up1 = Up(1024, 256) self.up2 = Up(512, 128) self.up3 = Up(256, 64) self.up4 = Up(128, 64) self.outc = nn.Conv2d(64, n_classes, 1) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) x = self.up1(x5, x4) x = self.up2(x, x3) x = self.up3(x, x2) x = self.up4(x, x1) return torch.sigmoid(self.outc(x))

3. 损失函数与评估指标

3.1 Dice 损失实现

Dice 系数更适合医学图像中不平衡目标的优化：

def dice_loss(pred, target, smooth=1.): pred = pred.view(-1) target = target.view(-1) intersection = (pred * target).sum() return 1 - (2. * intersection + smooth) / (pred.sum() + target.sum() + smooth)

3.2 混合损失函数

结合交叉熵和 Dice 损失的优点：

def bce_dice_loss(pred, target): bce = nn.BCELoss()(pred, target) dice = dice_loss(pred, target) return bce + dice

3.3 评估指标

实现论文使用的 Dice 系数评估：

def dice_coeff(pred, target): pred = (pred > 0.5).float() return 2 * (pred * target).sum() / (pred.sum() + target.sum() + 1e-6)

4. 训练流程与超参数优化

4.1 训练循环配置

设置关键超参数和训练流程：

import torch.optim as optim model = UNet().cuda() optimizer = optim.Adam(model.parameters(), lr=1e-4) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'max', patience=2) def train_epoch(model, loader, optimizer): model.train() total_loss = 0 for batch in loader: inputs, targets = batch inputs, targets = inputs.cuda(), targets.cuda() optimizer.zero_grad() outputs = model(inputs) loss = bce_dice_loss(outputs, targets) loss.backward() optimizer.step() total_loss += loss.item() return total_loss / len(loader)

4.2 验证与模型选择

监控验证集性能并保存最佳模型：

def evaluate(model, loader): model.eval() dice = 0 with torch.no_grad(): for batch in loader: inputs, targets = batch inputs, targets = inputs.cuda(), targets.cuda() outputs = model(inputs) dice += dice_coeff(outputs, targets).item() return dice / len(loader) best_dice = 0 for epoch in range(100): train_loss = train_epoch(model, train_loader, optimizer) val_dice = evaluate(model, val_loader) scheduler.step(val_dice) if val_dice > best_dice: best_dice = val_dice torch.save(model.state_dict(), 'best_model.pth') print(f'Epoch {epoch}: Loss={train_loss:.4f}, Dice={val_dice:.4f}')

4.3 关键超参数设置

复现论文结果的参数配置：

5. 结果分析与可视化

5.1 性能对比

在 ISBI 测试集上的表现：

5.2 预测结果可视化

使用 matplotlib 展示分割效果：

import matplotlib.pyplot as plt def plot_results(input, target, prediction): plt.figure(figsize=(12,4)) plt.subplot(131) plt.imshow(input[0,0].cpu(), cmap='gray') plt.title('Input') plt.subplot(132) plt.imshow(target[0,0].cpu(), cmap='gray') plt.title('Ground Truth') plt.subplot(133) plt.imshow(prediction[0,0].cpu() > 0.5, cmap='gray') plt.title('Prediction') plt.show()

5.3 实际应用建议

1. 对于小型数据集，建议冻结编码器部分权重
2. 当分割目标较小时，可调整 Dice 损失权重至 0.7-0.8
3. 输入尺寸需保持与训练时一致，可通过滑动窗口处理大图
4. 二分类任务输出通道设为1并使用sigmoid，多分类使用softmax