ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PyTorch与Unet实现MRI肝脏分割:跳跃连接、损失函数与训练避坑全指南

PyTorch与Unet实现MRI肝脏分割:跳跃连接、损失函数与训练避坑全指南 简介基于PyTorch与UNet的MRI肝脏图像分割项目面向计算机相关专业毕业生及需要医学影像算法实战的开发者可直接作为毕业设计、课程设计或期末大作业使用。压缩包涵盖完整源码、数据集与说明文档总计1070个文件其中1065张PNG图像样本用于模型训练与验证4个Python脚本实现数据加载、UNet网络构建、训练及预测流程1个README说明项目配置与运行方法整体约21.64MB下载后经过严格调试即可运行。已有871人学习/下载项目目录结构清晰能够帮助读者快速复现肝脏MRI分割流程掌握数据预处理、网络训练与评估的完整链路同时可基于自带数据集进一步扩展实验是图像分割方向较为稳妥的实战学习起点。1. 为什么毕业设计都在用 PyTorch Unet 做 MRI 肝脏分割肝脏是医学图像分割里最适合做毕业设计的器官尺寸大、边界相对固定、公开数据集多。用 PyTorch Unet 做 MRI 肝脏图像分割标准套路是把 3D 体数据切成 2D 切片逐片预测肝脏掩膜再拼回 3D 结果。这个方案训练成本低一张显卡就能跑完网上也能找到不少可复现的源码和数据集。但拿到“源码数据集”只是入场券。真正决定答辩分数的是数据预处理、损失函数组合和训练停止时机这三件事。如果你能解释清楚 Unet 的跳跃连接为什么适合小样本医学图像这篇毕业设计已经完成了一半。下面从结构到代码把整条流水线拆开讲。2. Unet 网络结构拆解跳跃连接为什么适合小样本医学图像2.1 对称骨架下采样提语义上采样还原分辨率Unet 的结构是一个对称的 U 形。左侧编码器由卷积块和 2×2 最大池化组成每经过一次池化特征图尺寸减半、通道数翻倍右侧解码器由转置卷积或双线性上采样加卷积组成逐步把尺寸恢复回去。以最常见的配置为例输入 256×256 的单通道切片经过四轮下采样变成 16×16 的特征图再经过四轮上采样回到 256×256。如果只有这个 U 形主干解码器想把细节恢复出来是非常困难的——下采样已经把边缘信息平均掉上采样又只是插值边缘和纹理基本要靠网络自己“脑补”。所以 Unet 的关键设计是跳跃连接编码器每一层下采样之前的特征图会被直接拼接到解码器对应层上。拼接让解码器同时拿到“高分辨率的浅层细节”和“低分辨率的语义信息”肝脏和腹壁、胃肠道之间只有几个像素的灰度差时这些信息恰恰是浅层特征里最丰富的。对比 DeepLab 和 PSPNet 这类带空洞卷积的分割网络Unet 在医学小样本场景下更稳的原因也在这里空洞卷积需要预训练权重兜底而 Unet 从零训练就能收敛。代码层面跳跃连接在 PyTorch 里就是一行torch.cat([enc_feat, dec_feat], dim1)。标准 Unet 实现里编码器和解码器每一层的通道是对应的拼接之后接一个卷积就能自动调整通道如果你把编码器换成 ResNet 之类的 Backbone跳跃连接传过来的通道数和解码器上采样的输出可能不一致这时候需要在拼接之前加一个 1×1 卷积把通道投影到统一维度这个细节很多魔改版代码都没写清楚。2.2 深度、通道数和输入尺寸先别堆参数我第一次跑 Unet 的时候吃过一个亏把网络深度加到 6 层初始通道提到 128结果显存直接爆掉而且 Dice 并没有变好。后来才意识到肝脏在 MRI 上的灰度分布相对均匀不需要极深的感受野参数堆上去只会增加过拟合风险。下面这张表是我在这个任务上常用的初始配置也是大多数开源 Unet 代码默认的取值参数常用范围我先从哪个开始输入尺寸256×256 / 512×512显存够就 512×512编码器初始通道32 / 6464下采样层数3 / 4 / 54Batch size4 / 8 / 168512 输入时减到 4学习率1e-5 ~ 3e-4Adam 1e-4深度加到 5 时感受野变大但特征图分辨率变低细小边界会进一步丢失。肝脏边界本身比较平滑4 层已经够用。真要提升精度先动输入尺寸——从 256 换到 512对边界 Dice 的提升通常比加深网络更明显代价是显存和训练时间。如果显存见底优先减 batch size 而不是减输入尺寸因为 BatchNorm 层的统计量在小 batch 下会抖动影响收敛稳定性。对于第一次跑一个 Unet 网络的读者我建议先用这套默认参数跑通再逐个动其中一个变量。改网络结构之前先跑出 baseline不然你根本不知道改动是变好还是变坏。这里还有一个容易忽略的点PyTorch 里 Unet 没有官方预训练权重网上的实现大多是从 ImageNet 分类模型里拆编码器过来Decoder 参数名可能对不上加载时需要strictFalse。2.3 预训练权重MRI 任务里不要直接照搬 ImageNet看到这里你可能想问能用 ImageNet 预训练权重吗能用但要接受一个事实MRI 的灰度分布和自然图像差距非常大预训练权重里的卷积核并不认识“肝脏组织在 T1 加权下是什么纹理”。直接用预训练编码器微调前期收敛可能快一点但如果数据集很小反而容易把特征空间带偏。我一般这样处理Encoder 加载预训练权重Decoder 随机初始化整体学习率放在 1e-4 以下。如果数据集只有几十例干脆不用预训练从零训练反而少一个变量。很多开源的 Unet 代码里有一个pretrainedTrue的开关跑毕设时不要无脑打开先对比一次开与不开在验证集上的 Dice 再决定。提示用strictFalse加载预训练权重之前先打印两边的 state_dict 的 key 对比确认不是静默丢弃了大部分参数。另外不要迷信那些针对 3D 医疗数据设计的预训练模型。它们输入的是体素块体素间距和切片方向都和你手上的 2D 切片不一样强行迁移只会引入一个黑匣子。Unet 在肝脏分割上本来就是个强 baseline把训练细节管好比换一个花哨的预训练模型有用得多。3. 把原始 MRI 整理成训练集NIfTI 读取、空片过滤和增强3.1 NIfTI 读取与归一化MRI 没有 HU 值先做百分位裁剪和 z-score拿到数据集先分清格式。常见的是 NIfTI 文件.nii.gz带一个同名的 mask 文件也有数据集已经预处理成 PNG 切片。如果给的是 NIfTI第一步是读取并统一归一化方式。import nibabel as nib import numpy as np from scipy.ndimage import zoom def load_volume(nifti_path, label_path): img nib.load(nifti_path).get_fdata(dtypenp.float32) lab nib.load(label_path).get_fdata(dtypenp.int16) # NIfTI 里常见 NaN 和极端值先清理掉 img np.nan_to_num(img, nan0.0, posinf0.0, neginf0.0) # 有 mask 时按 mask 区域内统计 1% 和 99% 分位裁剪极端亮/暗点 if (lab 0).any(): lo, hi np.percentile(img[lab 0], [1, 99]) else: lo, hi np.percentile(img, [1, 99]) img np.clip(img, lo, hi) # z-score 归一化让不同扫描的灰度分布对齐 img (img - img.mean()) / (img.std() 1e-8) return img, lab这段代码里有两处很多人会写错。第一MRI 信号强度没有 CT 那样的绝对物理单位亨氏单位所以不能照搬 CT 分割代码里的固定窗口裁剪比如np.clip(img, -200, 200)这种写法在 MRI 上会直接把大部分组织细节削掉。第二百分位裁剪的统计范围尽量用 mask 区域内的像素因为整张图里背景噪声占比高会拉低分位数的参考价值。z-score 归一化的作用是让多个病例的灰度分布对齐。同一个患者在同一台设备上扫两次灰度都可能不一样不做归一化模型会把设备差异当成语义差异验证集 Dice 会明显掉。切完切片之后每个切片还可以再按切片做一次 z-score但我习惯用体数据级别的归一化稳定性更好。3.2 从 3D 体数据切 2D 切片方向、空片过滤和 mask 重采样MRI 体数据是三维的而 2D Unet 一次只看一个切片。切之前先确认数据的方向大多数 NIfTI 数据的第三个维度是轴向从头顶到脚底沿第三维切片能得到我们熟悉的横断面。少数数据可能是冠状位或矢状位排列最稳妥的方法是先print(img.shape)并对着切片可视化确认不要假设所有数据集都是轴向。def extract_slices(img, lab, target_size(256, 256), min_ratio0.001): slices, masks [], [] # 假设数据已经是轴向排列沿 z 轴逐片切 for idx in range(img.shape[2]): slc img[:, :, idx] msk lab[:, :, idx] # 肝脏区域占比低于阈值的切片直接丢掉 if (msk 0).sum() / msk.size min_ratio: continue # 图像用线性插值掩膜用最近邻插值避免出现中间灰度 scale (target_size[0] / slc.shape[0], target_size[1] / slc.shape[1]) slc zoom(slc, scale, order1) msk zoom(msk, scale, order0).astype(np.uint8) slices.append(slc[np.newaxis, :, :]) masks.append(msk[np.newaxis, :, :]) return slices, masks两个参数值得说明。order1是线性插值用于灰度图order0是最近邻插值用于掩膜。如果用线性插值缩放掩膜肝脏边缘会出现 0.5、0.3 这样的中间值训练时损失函数会对这些像素产生梯度直接污染分割边界。min_ratio用来过滤空片和只有零星几个像素的伪阳性切片这个值不是越大越好——肝脏在体数据里本来就可能只占中间几十片过滤太狠会把有效样本全丢掉。切片方向要不要统一取决于训练效果。如果发现预测掩膜里肝脏呈竖条状那大概率是训练时切片方向混了比如一部分是轴向、一部分是冠状位。处理方式是在切完所有病例后对每张切片做一次主轴方向检查确保所有输入图像里的肝脏都近似横躺。3.3 数据增强翻转、旋转够用别盲目上弹性变形医学图像分割的数据增强和自然图像不一样核心原则是“不能改变解剖结构的拓扑”。水平翻转是安全的90° 旋转也是安全的因为肝脏在腹腔里的相对位置不会因为这些几何变换而改变。def augment(img, msk): # 水平翻转和垂直翻转 if np.random.rand() 0.5: img np.flip(img, axis1) msk np.flip(msk, axis1) if np.random.rand() 0.5: img np.flip(img, axis2) msk np.flip(msk, axis2) # 随机 90° 旋转不引入插值误差 k np.random.randint(0, 4) img np.rot90(img, k, axes(1, 2)) msk np.rot90(msk, k, axes(1, 2)) return img.copy(), msk.copy()这里故意不写随机角度旋转和弹性变形。随机角度旋转比如旋转 15°需要插值会引入灰度变化弹性变形在医疗影像增强里有专门文献支持但要控制形变幅度和网格间距调不好会把肝脏拉成不规则的形状。Unet 对刚性变换的鲁棒性已经足够好毕业设计阶段用翻转和 90° 旋转足够把精力留在更值得调的地方。还有一个在“unet 训练自己的数据集”时很常见的错误增强写在训练集上但验证集也跟着增强。验证集必须保持原始分布否则测出来的 Dice 不是真实水平。把增强逻辑和 Dataset 解耦训练时传augmentTrue验证时传False是最稳妥的用法。4. 用 PyTorch 训一个肝脏分割 Unet损失函数搭配与训练循环4.1 损失函数怎么选BCE、Dice 还是混合肝脏分割的核心难点是类别不平衡一张 512×512 的切片里肝脏区域可能只占 10%30%其余全是背景。如果用纯 BCE二元交叉熵网络会学出一个“全都预测为背景”的局部最优解Loss 看着在降Dice 却是 0。常见做法是用 Dice Loss 或 BCE Dice 混合。Dice Loss 直接优化分割任务的目标指标对前景占比不敏感BCE 的梯度更平稳两者互补。损失函数优点缺点适用场景BCE收敛平稳梯度连续背景占比大时被背景主导前景占比大于 30%Dice Loss直接优化目标指标天然处理不平衡前景极小时梯度跳变前景占比 10% 以下BCE Dice两者互补收敛和优化目标兼顾需要调节权重肝脏分割的最稳选择Focal Loss对难分样本加权超参数多容易调崩边界极其模糊时考虑我的默认组合是BCE Dice比例 1:1。先跑通再根据验证集结果微调。直接上 Focal Loss 的翻车概率很高参数稍微没调好训练里 Dice 直接卡在 0.5 以下不动。4.2 训练循环和模型保存策略Dice 不涨时该看什么下面是一个可以直接用的混合损失实现import torch import torch.nn as nn import torch.nn.functional as F class DiceBCELoss(nn.Module): def __init__(self, smooth1e-6): super().__init__() self.smooth smooth def forward(self, logits, targets): probs torch.sigmoid(logits) targets targets.float() # BCE 部分 bce F.binary_cross_entropy(probs, targets, reductionmean) # Dice 部分按 batch 内每个样本算再取平均 inter (probs * targets).sum(dim(2, 3)) union probs.sum(dim(2, 3)) targets.sum(dim(2, 3)) dice 1 - (2 * inter self.smooth) / (union self.smooth) return bce dice.mean()这里有个细节Dice 在 batch 维度上先对每张图单独算再取平均而不是把整个 batch 的像素合并算。后者会被某张肝脏特别大的图片主导导致小肝脏样本的梯度被淹没。smooth参数是为了防止分母为 0取1e-6足够。训练循环本身不长但有几个关键点def train_epoch(model, loader, optimizer, criterion, device): model.train() losses [] for imgs, masks in loader: imgs imgs.to(device) masks masks.to(device) optimizer.zero_grad() logits model(imgs) loss criterion(logits, masks) loss.backward() optimizer.step() losses.append(loss.item()) return np.mean(losses)训练之前先确认torch.cuda.is_available()。如果你的 pytorch 环境搭建用的是 CPU 版一张 512×512 的切片前向加反向可能要 35 秒数据集几百张图跑一轮就是半小时起步整个调参周期会被拖得很难受。CUDA 版能装就装装完先跑一次python -c import torch; print(torch.cuda.is_available())验证。模型保存策略上不要只存最后一个 epoch 的权重。常见做法是每个 epoch 在验证集上算一次 DiceDice 创新高就覆盖保存best_model.pth。这样即使后几个 epoch 过拟合了你手里还有一个最佳版本。如果验证集 Dice 连续 10 个 epoch 不涨就把学习率降到原来的 1/5 再训而不是直接放弃。4.3 推理和阈值选取测试时增强要不要开训练完成后推理流程是固定结构模型设成eval()关闭梯度逐切片预测再按阈值二值化。def predict_volume(model, volume, device, threshold0.5): model.eval() preds [] with torch.no_grad(): for slc in volume: # volume 里的每张切片都是 [1, H, W]补一个 batch 维 t torch.from_numpy(slc).unsqueeze(0).unsqueeze(0).to(device) logit model(t) prob torch.sigmoid(logit).cpu().numpy()[0, 0] preds.append((prob threshold).astype(np.uint8)) return preds阈值不一定用 0.5。在验证集上对每个候选阈值比如 0.3 到 0.7步长 0.05分别算 Dice取最优阈值用于测试集这一步通常能让最终分数提高 12 个点。如果训练时增强用了翻转和旋转测试时增强可以只做水平翻转把原图和翻转图分别预测取概率平均再二值化。代价是推理时间翻倍毕设阶段数据量不大开不开都行但要在报告里写清楚。5. 训练和验证中的避坑清单显存爆掉、Dice 不涨、掩膜偏移5.1 训练到一半显存溢出OOM现象训练在某个 epoch 中途直接报CUDA out of memory前面跑得好好的突然就爆了。原因分两类。一是输入尺寸和 batch size 设置过大512×512 的输入配 batch16单卡不可能装下二是训练过程中随着 epoch 增加PyTorch 的自动求导图把中间变量积累下来峰值显存比第一个 step 更高。解决先把 batch size 降到 4如果还爆再把输入尺寸降回 256×256。不要同时改两个变量否则无法定位是哪个参数造成的。PyTorch 的torch.cuda.amp.autocast()混合精度也能省 30%40% 显存改造代价很小。还有一个常用手段是梯度累积optimizer.zero_grad()改成每 4 个 batch 才执行一次optimizer.step()效果等价于 batch size 扩大 4 倍缺点是训练节奏变慢需要同步调整学习率。5.2 Loss 一直在降但验证集 Dice 卡在 0.7 以下现象训练集 Loss 从 0.8 降到 0.2但验证集 Dice 始终在 0.60.7 之间抖动怎么调都上不去。原因往往是输入数据和掩膜的预处理不一致训练时切片做了 z-score 归一化但验证或推理时直接用了原始灰度或者掩膜在缩放时用了线性插值导致模型学习的边界是模糊的。解决把训练、验证、推理三段代码里的预处理函数抽成同一个不要各写一份。检查掩膜 resample 是否用了order0以及预测后处理时有没有把概率图 resample 回原始分辨率和原始方向。还有就是前面提到的阈值没扫0.5 固定阈值可能恰好落在最差位置。5.3 预测掩膜整体偏移或方向不对现象单个切片看 Dice 很高把切片堆回 3D 体数据后发现肝脏整体往一个方向偏了几毫米或者上下颠倒。原因NIfTI 文件的坐标方向没有统一。有的数据是 RAS 方向有的是 LPS 方向体素原点也不一致直接按数组索引切片再重组几何位置全是错的。解决读取 NIfTI 时用nib.load(path).header查看qform和sform在数据预处理阶段把所有体数据对齐到同一个方向再做切片。简单起见先对每个病例的 mask 计算质心确认质心都在同一侧如果发现部分病例左右颠倒把这部分数据沿轴翻转后再进入训练集。这一步属于血泪经验——我见过有人训练了三天最后发现 20% 的数据方向反了重做预处理后 Dice 直接涨了 5 个点。5.4 训练集 Dice 高、验证集 Dice 崩现象训练集 Dice 0.92验证集只有 0.78差距明显。原因有两个一个是过拟合另一个是数据划分泄漏。有些代码把同一个病人的切片同时分到训练集和验证集导致验证集里出现了和训练样本几乎一样的切片测出来虚高而换个病人做验证分数立刻崩掉。解决数据划分必须按病例级别分割同一个病例的所有切片只能进训练集或验证集不能混着放。具体做法是先列出所有病例 ID按 8:2 或 7:3 划分再按病例 ID 取切片。过拟合时优先加数据增强、减小模型容量、把 Dropout 加在解码器末尾。5.5 多次训练结果不一致答辩无法复现现象同样的代码跑两次Dice 相差 3 个点以上最终用哪个结果都不踏实。原因PyTorch 默认的随机数种子没有固定数据加载时 DataLoader 的 shuffle 也随机加上 GPU 卷积的原子操作不确定性训练结果天然有波动。解决在训练脚本开头固定所有随机源import torch import numpy as np import random def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)同时把 DataLoader 的shuffle随机种子固定并在报告里注明最终用的是第几个 epoch 的权重。顺带把推理阈值也固定下来不然每次跑预测结果都会变答辩演示时最容易翻车。这个细节虽然不涨分但能直接决定你的实验部分有没有可信度。6. 把模型结果收进毕业设计切片可视化、Dice 评估和报告话术6.1 用边界轮廓图验证分割质量别只盯着 DiceDice 是一个笼统的数字两块形状完全不同的掩膜可能得到相同的 Dice。验证时除了算指标一定要把预测结果以可视化方式输出把预测掩膜边界叠加到原始 MRI 切片上同时用红色标出漏检区域、绿色标出误检区域。这样一眼就能看出模型是在边界处模糊还是在某些解剖结构处系统性地漏掉。生成对比图的代码很简单用 matplotlib 把原图、真实掩膜、预测掩膜按三列排开保存成一行一行的对比图集。答辩时放三组代表性切片效果最好、中等、最差比放十页曲线有说服力得多。6.2 把指标写进报告Dice、IoU 和边界距离怎么表述毕设报告里只写一个 Dice 分数不够。建议按病例维度统计 Dice、IoU、Precision、Recall 四列的均值和标准差再加一个表记录不同 epoch 的验证集指标变化。这样既能体现实验的完整性也能解释你最后选择的模型权重是哪一个。如果数据集包含多位患者的注释还可以按患者统计 Dice 的分布指出模型对哪一类图像比如脂肪肝、有肿瘤的肝脏表现偏弱。这部分内容放到讨论章节里直接成为你论文的亮点——大多数人只会贴一张训练 loss 曲线你拿出按病例粒度的误差分析答辩老师会认为你真的理解了模型。在我带过的毕设里最后得分高的并不是网络改得最花哨的那个而是数据流最清晰、能说清楚每一步为什么这么做的那个。指标计算和可视化脚本要单独整理成一个文件保证评审老师能直接复现你的数字。希望这篇笔记能帮你把 PyTorch Unet 这条流水线从头到尾理顺少走几周弯路。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进