简介:本资源面向医学图像分割方向的开发者与研究者,提供一套基于Transformer-Unet的超声腹部多器官语义分割完整方案,覆盖肝脏、肾脏、胰腺、血管、肾上腺、胆囊、骨骼、脾脏等目标类别。包内共1888个文件,以1852张png图像数据为主,辅以18个py脚本、15个pyc编译文件及txt说明与readme文档,压缩包约43.54MB,结构清晰便于直接上手。训练脚本采用AdamW优化器与余弦退火学习率衰减,损失函数为交叉熵,可输出训练集与验证集的loss、IoU曲线、学习率衰减曲线、训练日志及数据集可视化图像,并保存最优与最终权重;evaluate脚本用于计算测试集的IoU、召回率、精确率与像素准确率;predice脚本可生成gt及gt+image掩膜图像。代码注释详尽,README提供训练自有数据的傻瓜式流程,已有583人学习,适合希望快速复现并迁移到自身超声数据的中高级读者。
1. 超声腹部多器官分割:Transformer-Unet 这套代码到底能跑出什么
腹部超声的语义分割,做过的人都知道它和 CT、MRI 完全不是一个难度量级。超声图像本身信噪比低、器官边界模糊、还带着大量声影和斑点噪声,同一个肝脏在不同切面下灰度分布能差出一大截。传统 Unet 在这类数据上能跑,但遇到器官紧贴、边界被遮挡的情况,卷积核的局部感受野就有点力不从心。这套基于 Transformer-Unet 的超声腹部多器官分割代码,解决的正是这个问题:它把 Transformer 的全局注意力机制嵌进 Unet 的编解码结构里,让模型在低对比度的超声图上也能抓住器官之间的长距离空间关系。
代码覆盖了 liver、kidney、pancreas、vessels、adrenals、gallbladder、bones、spleen 等多个腹部器官类别,配套数据集、训练脚本、评估脚本、推理脚本一应俱全。优化器用 AdamW,学习率走余弦退火,损失函数是交叉熵。适合两类人:一是想直接拿一套能跑通的超声分割 baseline 做实验的算法工程师,二是想搞清楚 Transformer 和 Unet 怎么结合、又不想从零搭框架的开发者。下面按「资源是什么 → 怎么用 → 坑在哪」的顺序拆开讲。
2. Transformer-Unet 的结构选型:为什么不是纯 Transformer 也不是纯 Unet
2.1 编码器里 Transformer 块放在哪一层
纯 Transformer 做分割的问题在于它需要大量数据才能收敛,而医学超声的标注数据通常只有几百到几千张,直接上 ViT 类结构很容易过拟合。纯 Unet 的问题前面说了,局部卷积对长距离依赖建模不足。这套代码的思路是在 Unet 的编码器末端和瓶颈层引入 Transformer 块,浅层仍然用卷积提取局部纹理特征,深层才用自注意力捕捉全局语义关系。
具体来说,编码器前几级是标准的卷积下采样,到了分辨率降到 1/8 或 1/16 之后,特征图尺寸已经比较小,此时接入 Transformer 块的计算开销可以接受,同时注意力机制能在低分辨率上有效建模器官之间的空间位置关系。解码器部分保持 Unet 的跳连结构,把编码器浅层的高分辨率特征和解码器的语义特征拼接,保证边界分割的精细度。
这种混合结构的好处是:浅层卷积负责「看清楚纹理」,深层注意力负责「理清楚关系」。对于超声图像里肝脏和胆囊紧贴、肾脏和肾上腺边界模糊这类场景,全局注意力能提供卷积感受野覆盖不到的位置线索。
2.2 损失函数与优化器的搭配逻辑
代码里用的是交叉熵损失,优化器 AdamW,学习率余弦退火。这三者的组合不是随便选的。
交叉熵在多类别分割里是最直接的逐像素分类损失,它对每个像素的类别预测做 softmax 后计算负对数似然。超声分割里类别不平衡是常态——肝脏区域通常占图像面积很大,肾上腺、血管这些类别像素占比很小。交叉熵在这种情况下会被大类主导,小类分割效果差。常见做法是加类别权重或者换成 Dice + CE 的组合损失,代码里目前是纯交叉熵,如果你自己的数据里小类效果不好,这是第一个要改的地方。
AdamW 相比 Adam 的核心区别是权重衰减的处理方式。Adam 把 L2 正则加在梯度里,AdamW 把权重衰减直接作用在参数更新上,解耦了自适应学习率和正则化。在 Transformer 结构里,这种解耦对训练稳定性有实际帮助,尤其是注意力层的参数。
余弦退火的学习率策略是让学习率从初始值按余弦曲线下降到接近零,中间不做重启。它的好处是训练后期学习率足够小,模型能在损失曲面底部精细收敛,不会像阶梯衰减那样在后期突然跳变导致震荡。代码里初始学习率需要根据你的 batch size 调整,常见做法是 batch size 16 时初始 lr 设在 1e-4 到 3e-4 之间。
2.3 数据集组织与类别定义
数据集按器官类别分文件夹或按掩膜标签值区分,具体格式参考 README。从摘要描述看,类别包括 liver、kidney、pancreas、vessels、adrenals、gallbladder、bones、spleen 等。这里要注意的是,不同超声设备、不同扫查切面下同一器官的形态差异很大,如果数据集里混了多种设备的数据,建议先做一轮可视化检查,确认标注一致性。
提示:拿到数据集后先跑一遍可视化脚本,把图像和掩膜叠加看一眼。超声数据里标注错位、掩膜偏移的情况比 CT 数据常见得多,不先检查直接开训,loss 曲线会教你做人。
3. 训练脚本怎么跑:从数据加载到权重保存的完整链路
3.1 环境准备与依赖安装
代码是 Python 写的,依赖 PyTorch 生态。建议用 conda 建一个独立环境,避免和系统里的其他 PyTorch 版本冲突。
conda create -n us_seg python=3.8 -y conda activate us_seg pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install numpy opencv-python matplotlib tqdm tensorboard这里 PyTorch 版本根据你的 CUDA 版本选,cu118 对应 CUDA 11.8。如果没有 GPU,把 index-url 换成 CPU 版本即可,但训练会慢很多。tensorboard 是用来可视化训练曲线的,代码里 train 脚本会生成 loss、iou 曲线和学习率衰减曲线,用 tensorboard 看最方便。
3.2 train 脚本的参数配置与启动
train 脚本是整套代码的核心入口。它负责数据加载、模型初始化、训练循环、验证、保存权重和日志。启动前需要确认几个关键参数:
# config.py 或 train.py 里的参数段(示意,以实际代码为准) CFG = { "data_root": "./data/ultrasound", # 数据集根目录 "num_classes": 9, # 器官类别数,按你的数据集改 "img_size": 256, # 输入图像尺寸 "batch_size": 16, # 显存不够就降到 8 或 4 "epochs": 200, # 训练轮数 "lr": 3e-4, # 初始学习率 "weight_decay": 1e-4, # AdamW 的权重衰减 "optimizer": "AdamW", "scheduler": "CosineAnnealingLR", "loss": "CrossEntropyLoss", "save_dir": "./checkpoints", # 权重保存路径 }num_classes 一定要和你的数据集类别数对齐,多一个少一个都会导致训练时报维度错误。img_size 如果显存吃紧可以降到 128,但超声图像降太多会丢失小器官的边界信息,建议不低于 192。batch_size 和 lr 是联动的,batch 减半时 lr 也建议相应减小。
启动训练:
python train.py --config config.yaml --gpu 0训练过程中,脚本会在 save_dir 下生成训练日志、loss 曲线图、iou 曲线图、学习率衰减曲线,以及数据集可视化图像。权重文件会保存 best 和 last 两个,best 是按验证集 iou 最高时保存的,last 是最后一个 epoch 的。如果你中途中断了想继续训,看代码里有没有 resume 参数,没有的话需要手动加载权重。
3.3 训练日志里该盯什么
train 脚本输出的日志里,重点看三个东西:训练 loss 是否稳定下降、验证 iou 是否跟着涨、学习率曲线是否符合余弦形状。
如果训练 loss 下降但验证 iou 不涨,大概率是过拟合,考虑加数据增强或减小模型参数量。如果训练 loss 震荡严重,先检查学习率是不是太大,或者 batch size 太小导致梯度噪声大。如果学习率曲线不是平滑余弦而是阶梯状,检查 scheduler 的 T_max 设置是否和 epochs 匹配。
代码会生成数据集可视化图像,这个别跳过。它把输入图像、GT 掩膜、当前预测掩膜叠在一起显示,能直观看到模型在哪些器官上分得好、哪些分得差。超声图像里胆囊和血管经常被混淆,可视化一看便知。
4. 评估与推理:iou、recall、precision 怎么算,推理掩膜怎么生成
4.1 evaluate 脚本的指标计算逻辑
evaluate 脚本用于在测试集上评估模型。它加载训练好的 best 权重,对测试集逐张推理,然后计算 iou、recall、precision、像素准确率等指标。
python evaluate.py --weights ./checkpoints/best.pth --data_root ./data/ultrasound --split test这里要注意训练集、验证集、测试集的划分。训练集用于网络拟合,验证集用于调参和选 best 权重,测试集只在最后评估时用一次。如果你把测试集也拿去调参了,那评估结果就没有参考意义了。代码里 split 参数控制用哪个划分,常见做法是 7:1:2 或 8:1:1。
iou 的计算是逐类别的,每个器官单独算 iou 再取平均。recall 和 precision 也是逐像素统计后按类别汇总。像素准确率是全局的,所有像素里预测正确的比例。超声分割里像素准确率通常很高(因为背景和大器官占多数像素),所以别只看这一个指标,重点看小类器官的 iou。
4.2 predict 脚本生成推理掩膜
predict 脚本用于对单张或批量图像做推理,生成 GT 掩膜和 GT+image 叠加的掩膜图像。
python predict.py --weights ./checkpoints/best.pth --input ./samples --output ./results输出结果里会包含三类图:原始输入图像、GT 掩膜(如果输入带了标注)、预测掩膜叠加在原图上的效果图。叠加图是半透明掩膜盖在原图上,不同器官用不同颜色区分,方便直观检查分割效果。
推理时要注意输入图像的预处理要和训练时一致。训练时如果做了归一化、resize、通道转换,推理时也要做同样的操作。代码里如果封装了 transform 函数,直接复用即可,不要手动写一套不一样的预处理,否则推理结果会莫名其妙地差。
4.3 指标不达预期时的排查顺序
拿到评估结果后,如果 iou 明显低于预期,按这个顺序排查:
先看可视化结果,确认是整体分割差还是个别器官差。整体差通常是训练不充分或学习率不对,个别器官差通常是类别不平衡或标注质量问题。再看混淆矩阵(如果代码有输出的话),确认模型把哪些类别搞混了。最后看训练日志里的验证 iou 曲线,如果验证 iou 在某个 epoch 后开始下降,说明过拟合了,需要早停或加正则。
5. 避坑与常见问题:超声分割训练里那些翻车现场
5.1 现象:训练 loss 正常下降但验证 iou 始终在 0.3 以下
原因:最常见的是数据预处理不一致。训练时用了某种归一化(比如 ImageNet 均值方差),验证和推理时忘了做同样的归一化,导致输入分布偏移。另一个可能是类别标签从 0 开始还是从 1 开始没对齐,模型学的是错位的类别。
解决:检查 train 和 evaluate 脚本里的 transform 是否一致,检查数据集里掩膜的像素值范围和类别映射关系。打印几张训练图像和对应掩膜的像素值分布,确认没有异常值。
5.2 现象:训练到一半 loss 突然变成 nan
原因:学习率太大导致梯度爆炸,或者超声图像里有全黑或全白的异常样本,经过归一化后产生极端值。AdamW 虽然比 Adam 稳定,但初始 lr 设到 1e-3 以上仍然有风险。
解决:把初始学习率降到 1e-4 或 3e-4,加梯度裁剪(torch.nn.utils.clip_grad_norm_),检查数据集里有没有损坏的图像文件。如果用了混合精度训练,nan 的概率会更高,先关掉 AMP 跑一轮确认。
5.3 现象:小器官(肾上腺、血管)的 iou 接近零
原因:类别极度不平衡,交叉熵损失被大类主导,小类的梯度信号被淹没。超声图像里肾上腺和血管的像素占比可能不到 1%,模型倾向于全部预测为背景或大类。
解决:换损失函数,用 Dice Loss 或 CE + Dice 的组合损失。或者在交叉熵里加类别权重,给小类更高的权重。另一个办法是在数据采样时做重采样,让包含小类的图像更频繁地出现。
5.4 现象:推理结果比验证时差很多
原因:推理时的图像尺寸和训练时不一致,或者推理脚本里的模型加载方式有问题(比如没加载 best 而是加载了 last,或者权重 key 不匹配)。
解决:确认推理时的 img_size 和训练时一致,确认加载的权重文件是 best.pth,打印模型加载时的 missing keys 和 unexpected keys,如果有大量不匹配说明权重没加载对。
5.5 现象:训练速度异常慢,GPU 利用率低
原因:数据加载是瓶颈,num_workers 设得太小或者数据集读取方式效率低。超声图像如果是 PNG 格式且数量多,IO 可能成为瓶颈。
解决:把 num_workers 调到 4 或 8,用 pin_memory=True,如果数据集不大可以直接把图像预加载到内存里。另外检查有没有在训练循环里做了不必要的 CPU 计算,比如每步都画图或写日志。
6. 进阶技巧:用可视化反推模型到底学到了什么
训练跑通、指标也还行之后,下一步是搞清楚模型到底在看哪里。超声分割里,模型可能学到了「肝脏总是在图像某个区域」这种位置先验,而不是真正的边界特征。要验证这一点,可以做一个简单的消融:把输入图像的部分区域遮挡,看预测结果怎么变。
import torch import numpy as np def occlusion_test(model, image, mask_size=32, stride=16): """ 对输入图像做滑动遮挡,观察预测掩膜的变化。 image: (1, C, H, W) tensor mask_size: 遮挡块大小 stride: 滑动步长 """ model.eval() base_pred = model(image).argmax(dim=1) # 原始预测 H, W = image.shape[2], image.shape[3] heatmap = np.zeros((H, W)) for y in range(0, H - mask_size + 1, stride): for x in range(0, W - mask_size + 1, stride): occluded = image.clone() occluded[:, :, y:y+mask_size, x:x+mask_size] = 0 # 遮挡区域置零 with torch.no_grad(): pred = model(occluded).argmax(dim=1) # 统计预测变化的比例 diff = (pred != base_pred).float().mean().item() heatmap[y:y+mask_size, x:x+mask_size] = diff return heatmap这段代码的逻辑是:把图像分成小块,每次遮挡一块,看模型预测变化多大。变化大的区域说明模型依赖那里的信息做决策。如果遮挡肝脏区域导致预测大变,说明模型确实在看肝脏;如果遮挡背景区域也导致预测大变,说明模型可能学到了不该学的位置捷径。
heatmap 可以用 matplotlib 画出来叠加在原图上,直观看到模型的「注意力」分布。这个技巧在调试超声分割时特别有用,因为超声图像的很多区域本身就是噪声,模型如果依赖噪声区域做预测,泛化能力一定差。
另一个进阶方向是测试时增强(TTA)。对同一张测试图像做水平翻转、小角度旋转,分别推理后把结果平均,通常能涨 1 到 2 个点的 iou。代价是推理时间翻几倍,看你的场景能不能接受。
注意:TTA 对超声图像的效果取决于你的数据增强策略。如果训练时没做旋转增强,测试时做旋转 TTA 可能反而掉点。TTA 的变换要和训练增强对齐。
从那以后我每次跑新的分割任务,都强制先跑一遍可视化再开训,确认图像和掩膜对齐、类别映射正确、预处理一致,这三步不做完绝不动手调参。希望帮到你。
本文还有配套的精品资源,点击获取