UNet 医学图像分割实战：PyTorch 2.0 实现 4 层编码器-解码器结构

UNet 医学图像分割实战：PyTorch 2.0 实现 4 层编码器-解码器结构

医学图像分割是计算机视觉领域的重要应用方向，而UNet凭借其独特的U型结构和跳跃连接机制，已成为该领域的标杆模型。本文将带你从零开始实现一个标准的4层UNet网络，并深入探讨其在医学图像分割任务中的应用技巧。

1. UNet架构设计原理

UNet的核心思想是通过编码器-解码器结构实现多尺度特征融合。编码器负责提取图像的层次化特征，而解码器则逐步恢复空间分辨率。中间的跳跃连接（skip connection）机制能够将低层细节信息与高层语义信息相结合，显著提升分割精度。

关键组件对比表：

提示：PyTorch 2.0的torch.compile()可以显著提升UNet的推理速度，建议在模型初始化后立即调用

2. 模块化代码实现

我们先构建四个核心模块，采用面向对象的设计思想：

import torch import torch.nn as nn import torch.nn.functional as F 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): super().__init__() 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) class OutConv(nn.Module): """输出层卷积""" def __init__(self, in_channels, out_channels): super().__init__() self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1) def forward(self, x): return self.conv(x)

3. 完整UNet网络集成

基于上述模块，我们可以构建完整的4层UNet：

class UNet(nn.Module): def __init__(self, n_channels=1, n_classes=2): super(UNet, self).__init__() self.n_channels = n_channels self.n_classes = n_classes # 编码器路径 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, 1024) # 解码器路径 self.up1 = Up(1024, 512) self.up2 = Up(512, 256) self.up3 = Up(256, 128) self.up4 = Up(128, 64) self.outc = OutConv(64, n_classes) # PyTorch 2.0优化 self = torch.compile(self) 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 logits

网络参数统计：

model = UNet() total_params = sum(p.numel() for p in model.parameters()) print(f"总参数量：{total_params/1e6:.2f}M") # 约31.04M

4. 医学图像数据处理技巧

医学图像通常具有以下特点：

• 单通道灰度图像
• 高分辨率但样本量少
• 类别不平衡严重

数据增强策略：

from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.RandomRotation(15), transforms.ColorJitter( brightness=0.1, contrast=0.1), transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5]) ])

医学图像加载示例：

from torch.utils.data import Dataset import nibabel as nib # 用于读取NIfTI格式 class MedicalDataset(Dataset): def __init__(self, img_paths, mask_paths, transform=None): self.img_paths = img_paths self.mask_paths = mask_paths self.transform = transform def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img = nib.load(self.img_paths[idx]).get_fdata() mask = nib.load(self.mask_paths[idx]).get_fdata() if self.transform: img = self.transform(img) mask = self.transform(mask) return img, mask.long()

5. 训练优化与损失函数

医学图像分割需要特殊的损失函数处理类别不平衡：

def dice_loss(pred, target, smooth=1.): pred = pred.contiguous() target = target.contiguous() intersection = (pred * target).sum(dim=2).sum(dim=2) loss = (1 - ((2. * intersection + smooth) / (pred.sum(dim=2).sum(dim=2) + target.sum(dim=2).sum(dim=2) + smooth))) return loss.mean() class DiceBCELoss(nn.Module): def __init__(self, weight=None, size_average=True): super(DiceBCELoss, self).__init__() def forward(self, inputs, targets): # 二分类sigmoid inputs = torch.sigmoid(inputs) # 计算Dice损失 dice = dice_loss(inputs, targets) # 计算BCE损失 bce = F.binary_cross_entropy(inputs, targets.float()) return dice + bce

训练循环关键代码：

def train_epoch(model, loader, optimizer, criterion, device): model.train() running_loss = 0.0 for images, masks in loader: images = images.to(device) masks = masks.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, masks) loss.backward() optimizer.step() running_loss += loss.item() return running_loss / len(loader)

6. 模型评估与可视化

医学图像分割需要专业的评估指标：

def calculate_metrics(pred, target): pred = (pred > 0.5).float() target = target.float() tp = (pred * target).sum() fp = (pred * (1-target)).sum() fn = ((1-pred) * target).sum() precision = tp / (tp + fp + 1e-8) recall = tp / (tp + fn + 1e-8) dice = 2*tp / (2*tp + fp + fn + 1e-8) return precision, recall, dice

结果可视化函数：

import matplotlib.pyplot as plt def plot_results(image, mask, pred): plt.figure(figsize=(15,5)) plt.subplot(1,3,1) plt.imshow(image[0], cmap='gray') plt.title('Input Image') plt.subplot(1,3,2) plt.imshow(mask[0], cmap='gray') plt.title('Ground Truth') plt.subplot(1,3,3) plt.imshow(pred[0] > 0.5, cmap='gray') plt.title('Prediction') plt.show()

7. PyTorch 2.0特性优化

利用PyTorch 2.0的新特性可以显著提升性能：

# 混合精度训练 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() # 使用torch.compile加速 model = torch.compile(model, mode='max-autotune')

性能对比测试结果：