U-Net 图像分割实战：PyTorch 实现细胞边缘检测，Dice 系数达 0.95

2026/7/8 16:49:35

U-Net 图像分割实战：PyTorch 实现细胞边缘检测，Dice 系数达 0.95

医学图像分割一直是计算机视觉领域的重要研究方向，尤其在细胞分析、病理诊断等场景中，精确的边缘检测直接关系到后续分析的准确性。本文将带你从零实现一个基于PyTorch的U-Net模型，并在ISBI细胞分割数据集上达到0.95的Dice系数。不同于理论讲解，我们将聚焦工程实现中的关键细节，包括数据增强策略、损失函数选择、训练技巧等实战经验。

1. 环境准备与数据加载

首先需要配置合适的开发环境。推荐使用Python 3.8+和PyTorch 1.10+，这些版本在兼容性和性能上都有良好表现。以下是核心依赖的安装命令：

pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html pip install opencv-python scikit-image tqdm

ISBI数据集包含30张512x512的细胞显微镜图像，每张图都有对应的标注掩膜。我们需要自定义Dataset类来加载数据并进行预处理：

from torch.utils.data import Dataset import cv2 import os import glob import random class CellDataset(Dataset): def __init__(self, data_dir, transform=None): self.image_paths = sorted(glob.glob(f"{data_dir}/images/*.png")) self.mask_paths = sorted(glob.glob(f"{data_dir}/masks/*.png")) self.transform = transform def __len__(self): return len(self.image_paths) def __getitem__(self, idx): image = cv2.imread(self.image_paths[idx], cv2.IMREAD_GRAYSCALE) mask = cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE) # 归一化并增加通道维度 image = image[None, ...] / 255.0 mask = mask[None, ...] / 255.0 if self.transform: sample = self.transform({ 'image': image, 'mask': mask }) image, mask = sample['image'], sample['mask'] return image, mask

数据增强对医学图像尤为重要，因为标注数据通常稀缺。我们使用Albumentations库实现了一套针对细胞图像的增强策略：

import albumentations as A train_transform = A.Compose([ A.RandomRotate90(), A.Flip(), A.ElasticTransform(alpha=120, sigma=120*0.05, alpha_affine=120*0.03, p=0.5), A.GridDistortion(p=0.5), A.RandomBrightnessContrast(p=0.2), ])

2. U-Net模型架构实现

U-Net的核心在于其编码器-解码器结构和跳跃连接。我们采用模块化方式实现，每个组件都有清晰的职责：

import torch import torch.nn as nn class DoubleConv(nn.Module): """(卷积 => [BN] => ReLU) * 2""" def __init__(self, in_channels, out_channels): super().__init__() self.double_conv = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) def forward(self, x): return self.double_conv(x) class Down(nn.Module): """下采样：最大池化后接双卷积""" def __init__(self, in_channels, out_channels): super().__init__() self.maxpool_conv = nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_channels, out_channels) ) def forward(self, x): return self.maxpool_conv(x)

解码器部分需要处理特征图的上采样和跳跃连接。这里我们比较了转置卷积和双线性上采样两种方式：

上采样方式	参数量	计算成本	边缘效果
转置卷积	有	较高	较锐利
双线性插值	无	低	较平滑

实际实现中可以根据硬件条件选择：

class Up(nn.Module): """上采样模块""" def __init__(self, in_channels, out_channels, bilinear=True): super().__init__() if bilinear: self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) self.conv = DoubleConv(in_channels, out_channels) else: self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2) self.conv = DoubleConv(in_channels, out_channels) 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)

完整的U-Net模型将这些组件组合成对称结构：

class UNet(nn.Module): def __init__(self, n_channels=1, n_classes=1, bilinear=True): super(UNet, self).__init__() self.n_channels = n_channels self.n_classes = n_classes self.bilinear = bilinear 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, bilinear) self.up2 = Up(512, 128, bilinear) self.up3 = Up(256, 64, bilinear) self.up4 = Up(128, 64, bilinear) self.outc = nn.Conv2d(64, n_classes, kernel_size=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) logits = self.outc(x) return torch.sigmoid(logits)

3. 训练策略与损失函数

医学图像分割常用的损失函数有以下几种选择：

BCEWithLogitsLoss：标准的二分类交叉熵
Dice Loss：直接优化Dice系数
组合损失：结合BCE和Dice的优点

我们采用组合损失，在训练初期依赖BCE的稳定性，后期由Dice Loss主导提升分割精度：

class DiceBCELoss(nn.Module): def __init__(self, weight=0.5): super(DiceBCELoss, self).__init__() self.weight = weight def forward(self, inputs, targets): # BCE计算 bce = F.binary_cross_entropy(inputs, targets) # Dice系数计算 inputs = inputs.view(-1) targets = targets.view(-1) intersection = (inputs * targets).sum() dice = 1 - (2.*intersection + 1)/(inputs.sum() + targets.sum() + 1) return self.weight*bce + (1-self.weight)*dice

训练流程的关键参数配置：

def train_model(model, device, train_loader, val_loader, epochs=100): optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, 'max', patience=5, factor=0.5, verbose=True) criterion = DiceBCELoss(weight=0.7) best_dice = 0 for epoch in range(epochs): model.train() for images, masks in train_loader: images, masks = images.to(device), masks.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, masks) loss.backward() optimizer.step() # 验证阶段 val_dice = evaluate(model, val_loader, device) scheduler.step(val_dice) if val_dice > best_dice: best_dice = val_dice torch.save(model.state_dict(), 'best_model.pth')

提示：使用混合精度训练可以显著减少显存占用，允许更大的batch size。只需在训练代码中添加几行：
scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs = model(images) loss = criterion(outputs, masks) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

4. 评估与结果可视化

Dice系数（F1分数）和IoU是评估分割精度的两个主要指标：

def dice_coeff(pred, target): smooth = 1. pred_flat = pred.view(-1) target_flat = target.view(-1) intersection = (pred_flat * target_flat).sum() return (2. * intersection + smooth) / (pred_flat.sum() + target_flat.sum() + smooth) def iou(pred, target): intersection = (pred & target).float().sum() union = (pred | target).float().sum() return intersection / (union + 1e-6)

结果可视化可以帮助我们直观理解模型的表现。下图展示了在验证集上的预测结果：