PyTorch实现UNet进行MRI肝脏分割的完整实践指南
发布时间:2026/9/15 4:56:36来源:尧图网络
简介基于PyTorch与UNet的MRI肝脏图像分割完整工程适合毕业设计、课程设计及期末大作业场景也方便刚接触深度学习的同学快速上手。压缩包共1070个文件其中1065张PNG图像为肝脏MRI数据集与分割结果可视化4个Python脚本覆盖模型训练、验证与预测流程另有1个Markdown文档提供项目说明与使用指引整体大小21.82MB部署门槛低。已有472人学习浏览属于导师认可的高分项目。代码附有详细注释下载后按文档配置环境即可运行还能根据README调整网络结构与训练参数适合在此基础上扩展消融实验或界面展示完成度较高。1. 用PyTorch复现UNet做MRI肝脏分割项目里的每一步怎么落地UNet从2015年提出到现在快十年医学图像分割比赛里已经很少作为冠军方案出现但真要把一个MRI肝脏分割任务落地跑通我依然会先把UNet作为第一版基线。肝脏在MRI里的灰度分布和周围组织对比度不稳定边界模糊前景区域在整张切片里的占比又小这类任务最怕什么标注数据不多、模型又大又复杂。UNet的跳跃连接在这种场景下特别契合既保留深层语义又兜住边缘细节。整个项目包含PyTorch源码、MRI切片与掩码数据集、配套文档代码内有注释下载后简单部署就能开始训练。下面的内容按数据处理、网络实现、训练验证的顺序展开新手能照着跑有经验的可以重点看排错和训练技巧部分。2. MRI肝脏数据集解析与预处理流程2.1 数据集的构成与掩码读取思路项目里提供的是PNG格式的MRI切片文件命名形如 Aug_No_485.png。这类命名常见于批量增广后保存的结果Aug 前缀表示经过增广No_ 后面的数字对应当前样本在原始数据集里的编号。理解命名规则很重要因为掩码文件通常会沿用同样的编号前缀再做区分。常见做法是掩码文件叫 Aug_No_485_mask.png或放在独立的 masks 目录下项目根目录的README文档一般也会写明切片和掩码的对应关系。拿到程序包先花五分钟确认图像和掩码是否一一对应文件数量不一致会直接造成训练时索引错位。读取切片和掩码时要注意掩码的通道。MRI切片本身是灰度信息保存成PNG后是无符号8位整数。掩码通常是单通道前景肝脏区域像素值为255背景为0。用OpenCV读取后要立刻把掩码压到0/1范围这一步是后续损失函数计算的基础。如果掩码保留255而不做处理sigmoid输出范围是0到1跟255相差好几个量级梯度回传时会剧烈震荡。import cv2 # 常见命名约定图像与掩码同名多一个 _mask 后缀 image_path Aug_No_485.png mask_path Aug_No_485_mask.png image cv2.imread(image_path) # BGR 三通道 mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # 单通道灰度 mask[mask 0] 1 # 二值化到 0/1这里 image 是 HxWxC 的BGR数组mask 是 HxW 的单通道数组。第5行用掩码索引把所有非零像素设成1这样掩码取值就只有0或1。后续不管用BCELoss还是DiceLoss输入预测概率和掩码都处于同一个数值尺度损失不会因为像素值量纲不同而产生异常梯度。如果掩码里还含有标注为其他器官的像素也要在这一步一并过滤只保留肝脏对应区域。2.2 尺寸统一与Z-score归一化数据集里的MRI切片原始分辨率未必一致有的采集自256x256有的经过重建存成512x512。网络输入必须统一这里比较合适的大小是256x256。理由很直接UNet编码器做四次下采样256输入在瓶颈层得到16x16的特征图空间信息足够输入变成512后显存占用会翻好几倍单卡环境很容易撑爆。灰度归一化采用Z-score而不是直接除以255。医学影像的对比度范围不是0-255均匀分布的MRI的灰度值不具备绝对物理意义不同患者、不同序列之间的亮度范围差异很大。把整个训练集的像素统计出来做标准化能把不同切片拉到近似的分布区间。统计量只在训练集上计算验证集和测试集沿用训练集算出的均值方差。import numpy as np train_gray [] # 依次存放读入的训练集灰度图 pixels np.concatenate([img.ravel() for img in train_gray]) mean, std pixels.mean(), pixels.std() def normalize(img): # 加 1e-8 防止纯背景切片的标准差为 0 return (img - mean) / (std 1e-8)这里 np.ravel 把二维数组拉成一维pixels.mean() 和 pixels.std() 的统计范围是所有训练切片的所有像素而不是单张切片。normalize 函数在验证时也要复用不能因为测试集像素分布略有偏差就重新计算否则训练和评估的输入分布不一致验证指标会有虚高或虚低的问题。加 1e-8 是为了避免某张切片完全退化为同一灰度时标准差为0导致除零错误。2.3 数据增广策略与DataLoader封装Aug_ 前缀说明数据准备阶段已经做过一轮离线增广但训练时仍建议在DataLoader里再做在线增广。离线增广生成的样本数量固定在线增广每次epoch能生成不同的变换组合相当于变相扩大数据集规模。对于肝脏MRI我常用的在线增广组合是这样一组增广方式参数范围使用建议水平翻转概率0.5建议开启肝脏左右镜像不影响语义垂直翻转概率0.5建议关闭腹部MRI头足方向有解剖含义随机旋转±15度可开启超过30度会破坏肝脏轮廓亮度/对比度扰动系数0.8~1.2模拟不同扫描参数带来的灰度偏移弹性形变sigma3~5增强对呼吸伪影的鲁棒性垂直翻转这一行需要多说一句。网络在训练时很容易把位置先验当成主要特征垂直翻转后肝脏位置上下颠倒网络如果学到了“灰度区域在下半部分就是肝”这种错误关联验证时反而会因为在正常体位扫描的切片上表现不稳定而掉点。PyTorch的Dataset在取样本时必须保证图像和掩码应用相同的随机增广否则训练样本里会出现“图像翻转了而掩码没有翻”的错误标注。用同一个随机种子可以做到同步import torch from torch.utils.data import Dataset class LiverDataset(Dataset): def __init__(self, image_paths, mask_paths, size256): self.image_paths image_paths self.mask_paths mask_paths self.size size def __len__(self): return len(self.image_paths) def __getitem__(self, idx): # load_image / load_mask 内部完成读取、resize、归一化并返回 torch.Tensor img load_image(self.image_paths[idx], self.size) mask load_mask(self.mask_paths[idx], self.size) # 同步随机变换手动设置同一个随机种子 seed torch.randint(0, 2**32, (1,)).item() torch.manual_seed(seed) if torch.rand(1).item() 0.5: img torch.flip(img, dims[2]) mask torch.flip(mask, dims[2]) return img.float(), mask.float()关键在第12行的 torch.manual_seed(seed)紧接着的 torch.rand(1) 决定这一对样本要不要翻转。如果后续还有旋转、裁剪也必须在同一个种子下生成参数并分别应用到 img 和 mask。这里 img 的 shape 是 [3, H, W]mask 是 [1, H, W]垂直和水平方向对应 dims 里的 1 和 2代码里翻转的是 dims[2]即宽度维只翻转左右方向。提示image 和 mask 的同步变换是这个项目里最容易出错的环节。建议在训练前单独写一个循环遍历 Dataset检查每次返回的 mask 和 image 是否同步翻转再进入正式训练。3. UNet网络实现与训练配置3.1 编码器-解码器的结构逻辑UNet网络结构图的经典构成是U形左右对称左侧编码器右侧解码器底部是瓶颈层。编码器通过池化逐层降低分辨率同时增加通道数来抽象语义信息解码器通过上采样逐层恢复分辨率同时减少通道数。跳跃连接把编码器同层的输出拼接到解码器对应层的输入上让边界定位信息在深层传播时不至于丢失。这里的项目采用UNet论文里的标准通道配置也就是64-128-256-512-1024。如果显存充足可以把首层通道改成32来减少参数但32通道下采样到深层后特征表达能力会明显下降。医学图像纹理细节本就比自然图像少再把通道压缩网络会更难区分肝脏与周围组织。UNet阶段输入→输出通道特征图尺寸作用编码器第一层3 → 64256×256提取低层边缘信息编码器第二层64 → 128128×128中等纹理特征编码器第三层128 → 25664×64较抽象的器官区域特征编码器第四层256 → 51232×32深度语义特征瓶颈层512 → 102416×16全局上下文关系3.2 跳跃连接与DoubleConv实现UNet代码实现里每个卷积阶段都用DoubleConv也就是两次“3x3卷积BNReLU”。BatchNorm在医学图像任务里的效果取决于batch大小这个项目默认batch size取8BN统计量足够稳定。实现时要给卷积去掉bias因为BatchNorm自带平移和缩放卷积的bias是冗余参数。import torch.nn as nn class DoubleConv(nn.Module): 两次卷积 BatchNorm ReLU def __init__(self, in_c, out_c): super().__init__() self.block nn.Sequential( nn.Conv2d(in_c, out_c, 3, padding1, biasFalse), nn.BatchNorm2d(out_c), nn.ReLU(inplaceTrue), nn.Conv2d(out_c, out_c, 3, padding1, biasFalse), nn.BatchNorm2d(out_c), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.block(x)Conv2d第一层把 in_c 变成 out_c第二层保持 out_c 不变。padding1 的作用是让3x3卷积在stride为1时保持特征图尺寸不变。biasFalse 配合BN使用BN完成平移归一化后卷积bias没有意义去掉还能减少参数量、降低过拟合风险。UNet主体forward里跳跃连接的拼接口用 torch.cat 完成。以解码器第二次上采样为例输入来自上一层上采样后的128通道特征和编码器第二层输出的128通道特征拼接后是256通道所以解码层DoubleConv的输入通道写256。class UNet(nn.Module): def __init__(self): super(UNet, self).__init__() self.enc1 DoubleConv(3, 64) self.enc2 DoubleConv(64, 128) self.enc3 DoubleConv(128, 256) self.enc4 DoubleConv(256, 512) self.pool nn.MaxPool2d(2) self.bottleneck DoubleConv(512, 1024) self.up4 nn.ConvTranspose2d(1024, 512, 2, stride2) self.dec4 DoubleConv(1024, 512) self.up3 nn.ConvTranspose2d(512, 256, 2, stride2) self.dec3 DoubleConv(512, 256) self.up2 nn.ConvTranspose2d(256, 128, 2, stride2) self.dec2 DoubleConv(256, 128) self.up1 nn.ConvTranspose2d(128, 64, 2, stride2) self.dec1 DoubleConv(128, 64) self.out nn.Conv2d(64, 1, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) b self.bottleneck(self.pool(e4)) d4 self.dec4(torch.cat([self.up4(b), e4], dim1)) d3 self.dec3(torch.cat([self.up3(d4), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return torch.sigmoid(self.out(d1)).squeeze(1)forward 最后一步的 squeeze(1) 把输出从 [batch, 1, H, W] 压成 [batch, H, W]后面计算损失时掩码的shape也是 [batch, H, W]两边直接对齐。sigmoid 放在模型内部而不是损失函数里意味着损失函数统一用BCELoss接收概率如果改用BCEWithLogitsLoss就要把这里的 sigmoid 去掉两者不能混用否则会出现梯度传两次的问题。3.3 损失函数与DiceLoss的加权组合肝脏区域在整张MRI切片里占比通常低于10%直接用BCELoss会让网络倾向做全背景预测。Dice损失对前景与背景比例不敏感是医学分割里的经典做法。将两者按0.5:0.5加权相加既保留BCE对概率的校准能力又发挥Dice对目标区域重叠度的直接优化。def dice_loss(pred, target, smooth1e-6): pred pred.contiguous().view(-1) target target.contiguous().view(-1) intersection (pred * target).sum() return 1 - (2 * intersection smooth) / (pred.sum() target.sum() smooth) def combined_loss(pred, target): bce nn.BCELoss()(pred, target) dice dice_loss(pred, target) return bce * 0.5 dice * 0.5dice_loss 里 pred 和 target 各自拉平到一维intersection 就是逐元素相乘后的和。smooth1e-6 在肝脏这种前景占比小的任务里很有必要如果某张切片恰好全为背景pred.sum() 和 target.sum() 同时为0时加平滑项可以避免出现 nan。BCELoss 要求 pred 是 sigmoid 之后的概率而不是 logits这与 UNet forward 最后输出 torch.sigmoid 的结果保持一致。注意如果训练过程中损失出现 nan先检查 pred 是否已通过 sigmoid再检查 smooth 是否遗漏。这两处是分割任务损失函数最常见的两种踩坑点。4. 训练循环与验证指标4.1 超参数设置的默认起点这个项目的数据量在几百张切片这个数量级训练配置建议从下面这组超参数开始。这套参数在多个类似肝脏分割的数据集上都能稳定收敛不需要一开始做复杂搜索。超参数推荐值说明输入尺寸256×256UNet四次下采样后瓶颈层16×16batch size8约4GB显存低于4时BatchNorm不稳定优化器Adam自带动量修正适合医学分割学习率1e-4比分类任务小一个量级避免震荡训练轮数100小数据集配合早停够用学习率调度ReduceLROnPlateaupatience10factor0.5Adam 在这种任务上比 SGD 收敛更快但从 1e-3 起手往往前几个 epoch 损失下降不稳定最终 Dice 反而偏低。分段常数衰减也能用但固定在第50和第80个epoch衰减未必对应模型真正需要调参的时机。ReduceLROnPlateau 的策略是监控验证集 Dice连续10个epoch没有提升就把学习率减半相对更省心。batch size 如果因为显存限制降到4以下建议把BatchNorm换掉详见第5章。4.2 Dice系数与IoU的计算细节验证阶段计算Dice系数之前预测概率需要先做二值化。这里有一个容易出错的地方直接用sigmoid输出的概率图算Dice会得到一个高于实际水平的分数因为概率在0.4-0.6之间的模糊像素会被当作部分贡献计入重叠。正确做法是加一个二值化步骤阈值取0.5。def compute_dice(pred_mask, gt_mask, eps1e-6): pred (pred_mask 0.5).float() gt (gt_mask 0.5).float() intersection (pred * gt).sum() dice (2 * intersection eps) / (pred.sum() gt.sum() eps) return dice.item()IoU 和 Dice 有换算关系Dice 2 * IoU / (IoU 1)。Dice 达到0.85时对应的 IoU 大约是0.74这在肝脏分割上属于正常范围。有的论文只报Dice有的只报IoU毕业设计里建议两个指标都打印答辩时老师从哪个指标切入提问都能接住。两个指标都用 eps 避免除零。4.3 训练循环中的可视化观察点训练循环里只打印 loss 是不够的。loss 下降不代表分割形状正确网络有可能学到输出一个小圆斑来匹配肝脏区域的重心loss 同样会下降。我习惯每5个 epoch 保存一批验证集的预测可视化至少覆盖4个样本如果环境支持 tensorboard用 add_images 记录最方便。from torch.utils.tensorboard import SummaryWriter writer SummaryWriter(runs/liver_seg) for epoch in range(epochs): # 训练循环... if epoch % 5 0: val_pred predict(model, val_loader) # [N, 1, H, W] 概率图 val_mask get_val_mask(val_loader) # [N, 1, H, W] 掩码 writer.add_images(val_pred, val_pred[:4], global_stepepoch) writer.add_images(val_mask, val_mask[:4], global_stepepoch) writer.add_scalar(val_dice, val_dice, global_stepepoch)tensorboard 能同时看到标量和图像两条线索。如果 Dice 一直停在0.6从叠加图能一眼看出是整块漏检还是边缘模糊漏检要找损失函数和类别权重方向边缘模糊则要检查输入分辨率、解码器上采样是否充分或者下采样次数是否过多导致边缘信息恢复不足。仅看数值给不出这些线索。5. 常见坑与调优技巧5.1 类别不平衡带来的“全背景陷阱”训练日志里如果出现 accuracy 很高、Dice 却为0的情况基本可以断定模型退化成全背景预测。肝脏 MRI 里最常见的原因就是前景占比太小训练初期 BCE 的梯度压过了 Dice 的梯度把网络推向全零输出。解决办法是前20个 epoch 内把损失权重调成 bce0.3、dice0.7同时每5个 epoch 检查一次预测图中非零像素的占比确认模型没有退化。5.2 显存不足时的BatchNorm替换显存不够时先降输入尺寸不要一上来改网络结构。把256改成224显存大约省三分之一。如果显存占用远高于UNet常规水平先检查训练循环里是否保存了多个 batch 的中间变量。batch size 确实小到4以下时把 BatchNorm 换成 GroupNorm 是更稳妥的解法def replace_bn_with_gn(module, num_groups8): for name, child in module.named_children(): if isinstance(child, nn.BatchNorm2d): setattr(module, name, nn.GroupNorm(num_groups, child.num_features)) else: replace_bn_with_gn(child, num_groups)这段递归代码从根模块出发把模型里每一处 BatchNorm2d 找出来换成 GroupNorm。num_groups 固定为8不管通道数是64还是512都能整除。替换要在模型初始化之后执行因为 GroupNorm 需要 num_features 参数直接读原 BatchNorm 的属性即可。做完替换之后训练和推理都用同一套网络不需要额外改其他代码。5.3 归一化统计一致性与水平翻转TTA归一化的均值方差只能从训练集计算验证集和测试集都要复用同一组统计值。如果整个数据集一起算统计量验证集信息被提前用在预处理里验证分数会虚高。另外预测阶段 resize 的插值方式必须和训练一致训练用双线性插值缩放到256x256预测时也要用双线性否则分割边界会产生一到两个像素的偏移Dice 略微下降。训练完成后可以用水平翻转TTA再提一点分数。对同一张切片翻转后各预测一次把两次概率平均再二值化通常比单次推理多出0.5到1个百分点的 Dicedef tta_predict(model, img): model.eval() with torch.no_grad(): p1 model(img) p2 model(torch.flip(img, dims[3])) p2 torch.flip(p2, dims[3]) return (p1 p2) / 2.0dims[3] 表示在 NCHW 张量的宽度维度上翻转翻转后再预测一次并把结果翻回来最后与未翻转的预测取平均。TTA 依赖训练时是否启用了水平翻转增广如果训练样本里从未出现翻转TTA 反而会引入方向不一致的偏差所以在 2.3 节把水平翻转开启是这一步的前提。本文还有配套的精品资源点击获取
网站建设高端定制企业官网