TransUnet实战:腹部CT多脏器分割数据与训练全攻略

📅 发布时间:2026/9/27 23:08:57
TransUnet实战:腹部CT多脏器分割数据与训练全攻略
简介基于TransUNet的腹部多脏器分割实战资源配套完整代码、数据集与训练结果面向医学图像分割领域的学生、科研人员及工程师覆盖肝脏、左肾、右肾、脾脏以及背景等五类组织的自动分割任务可直接用于算法对比、毕业设计或课题预研。压缩包共包含1031个文件整体约200.83MB其中以986张PNG图像为主涵盖腹部CT原图与对应掩膜另有18个Python脚本、2个PTH权重文件、说明文档及训练日志等分别承担训练、验证、预测和配置说明体系完整且便于按需取用。已有536人学习下载。代码注释详细训练脚本会自动保存loss与IOU曲线、学习率衰减曲线、数据集可视化样本以及最佳/最终权重验证脚本输出测试集IOU、召回率、精确率、像素准确率等指标预测脚本生成GT与叠加掩膜图像。按README操作即可训练自定义数据项目迭代100个epoch后测试集像素准确率为0.986平均IOU为0.779适合直接复现或二次开发。1. 用 TransUnet 做腹部多脏器分割你可能低估了数据这一关想用 TransUnet 在腹部 CT 上同时分割肝脏、脾脏、双肾和胰腺不是从 GitHub 上拉一段代码就能直接跑出结果的。我第一次做这个方向时训练脚本只写了半天后面两周全在和小标签错位、类别不平衡、学习率乱跳较劲。这也是我说“实战”比“模型结构”更重要的原因代码、数据集、训练结果三者缺一不可。这篇笔记适合医学图像算法工程师、研究生以及想在自己私有的腹部 CT 数据上快速验证 Transformer 分割方案的人。新手最容易低估预处理熟手则容易忽略验证集的噪声两条路都会翻车。2. TransUnet 的结构与选型为什么腹部 CT 需要全局上下文2.1 腹部多脏器分割的难点边界模糊、小器官和类不平衡腹部 CT 里肝脏、脾脏、肾脏、胰腺都属于软组织灰度范围接近很多边界在 CT 值上根本没有清晰的跳变。胰腺是最典型的例子它周围是十二指肠和肠系膜在普通窗宽窗位下几乎融在一起纯卷积网络如果只看局部感受野很容易把胰腺误判成肠道。再加上呼吸蠕动带来的器官形变同一个器官在不同病例里的位置变化也很大。小器官带来的第二个问题是样本量失衡。一张 224×224 的切片里背景能占 90% 以上胰腺可能只有几十个像素。普通交叉熵会把大量梯度花在“正确预测背景”这件事上结果肝脏、脾脏学得不错胰腺却迟迟不启动。很多语义分割算法在自然图像上效果好换到腹部 CT 上立刻失灵就是这个原因。所以能不能把器官之间的相对位置关系利用起来就成了提高分割精度的关键。比如脾脏一定在左侧胆囊贴着肝脏胰腺在胃和主动脉之间。这种全局上下文正是 TransUnet 和纯 U-Net 拉开差距的地方。我经验是不用全局模型时胰腺这类边界模糊器官最容易碎成小片换到 TransUnet 之后至少整体形状是完整成团的。2.2 TransUnet 的架构CNN 编码 Transformer 全局建模 U-Net 解码TransUnet 的常见实现可以拆成三段看。第一段是 CNN 编码器通常用 ResNet-50负责把原始切片降采样成高语义特征图第二段把这个特征图切成固定大小的 patch展平成 token 序列后送入 Transformer encoder让每一个位置都能看到全图第三段是 U-Net 风格的解码器把 Transformer 输出的序列恢复到空间结构并通过跳跃连接把 CNN 编码器早期的高分辨率特征融合回来。我一般关注三个可调参数输入分辨率、patch 大小和 Transformer 层数。输入分辨率增大时token 数量会平方级增加显存会迅速吃紧Transformer 层数增加会拉长训练时间但对最终 Dice 的增益往往不是线性的。很多人把 TransUnet 当黑匣子其实真正需要动的只有这几个位置。下面是组件和常见参数的对照表方便你排查代码里哪个部分不合预期。模块作用常见参数CNN 编码器提取边缘、纹理和语义特征ResNet-50输入通道 1Patch embedding把 1/16 特征图切成 token 序列patch_size1 或 16hidden_dim768Transformer encoder全局建模器官间相对位置n_layers12n_heads12Decoder Skip逐级上采样并融合细节每级上采样 2 倍2.3 和 U-Net、Swin-Unet 的选型对比如果拿不定主意可以先把 U-Net、TransUnet、Swin-Unet 放在一起比较。U-Net 实现简单、显存占用低但长距离信息只能靠堆叠卷积层胰腺这类器官容易学成碎片Swin-Unet 有层级 Transformer理论上更强但训练收敛慢对学习率和数据增强都比 TransUnet 敏感不适合作为第一个基线。以我自己跑过的经验TransUnet 是单卡实验室最容易复现的选择。它需要的显存比 Swin-Unet 低解码器结构又是经典的 U-Net 形态资源不够时可以只降 Transformer 层数。选型步骤也不复杂先用 TransUnet 跑一个基线看哪个器官拖后腿再针对性地调整损失或输入如果显存实在不够再降 token 数量或换轻量解码器。模型全局建模能力显存占用小器官表现实现成熟度U-Net弱依赖卷积堆叠低一般边界模糊处易漏极高TransUnet强Transformer 编码中较好胰腺、胆囊有提升较高Swin-Unet强且有层级结构中高好但调参门槛高较新3. 数据集准备把腹部 CT 扫描和标注变成 TransUnet 能直接训练的样子3.1 公开数据集选型Synapse/BTCV 是常见起点腹部多脏器分割这条赛道里公开实验用得最多的是 Synapse/BTCV 多器官数据集包含 30 例腹部 CT标注了脾脏、右肾、左肾、胆囊、食管、肝脏、胃、主动脉、胰腺共 9 个结构。常见做法是 18 例训练、12 例验证。这个拆分不算大但对 TransUnet 这种带 Transformer encoder 的模型来说足够训练到可用的水平。如果你用的是公司内部数据集第一件事不是改模型而是统一标注规则。比如左右肾是否分开、胆囊是否包含内容物、主动脉是从腹腔干到髂分叉还是只取中间一段。这些规则直接影响类别数和边界质量也比网络结构更影响最终分数。我这里用 Synapse 的标签顺序来举例1 脾脏、2 右肾、3 左肾、4 胆囊、5 食管、6 肝脏、7 胃、8 主动脉、9 胰腺。3.2 预处理流程重采样、窗宽窗位、归一化腹部 CT 原始图像的分辨率和体素间距在不同医院、不同设备下差异很大。我一般会先把所有病例重采样到统一的间距例如轴内 1mm×1mm层厚 5mm。z 方向用 5mm 是为了减少切片数、降低训练显存如果你的小器官在 z 轴很薄比如胆囊还是用 3mm 更稳。下面是常用的预处理片段。import SimpleITK as sitk import numpy as np def preprocess_ct(image_path, label_path, target_spacing(1.0, 1.0, 5.0)): # 读取原始 CT 和标注 img sitk.ReadImage(image_path) lab sitk.ReadImage(label_path) # 重采样图像使用线性插值 resampler_img sitk.ResampleImageFilter() resampler_img.SetOutputSpacing(target_spacing) resampler_img.SetSize([int(sz * sp / ts) for sz, sp, ts in zip(img.GetSize(), img.GetSpacing(), target_spacing)]) resampler_img.SetInterpolator(sitk.sitkLinear) resampled_img resampler_img.Execute(img) # 重采样标签必须使用最近邻插值否则会产生新标签值 resampler_lab sitk.ResampleImageFilter() resampler_lab.SetOutputSpacing(target_spacing) resampler_lab.SetSize(resampled_img.GetSize()) resampler_lab.SetInterpolator(sitk.sitkNearestNeighbor) resampled_lab resampler_lab.Execute(lab) # 转 numpy并应用腹部软组织窗 image_np sitk.GetArrayFromImage(resampled_img).astype(np.float32) label_np sitk.GetArrayFromImage(resampled_lab).astype(np.uint8) image_np np.clip(image_np, -150.0, 250.0) image_np (image_np - (-150.0)) / (250.0 - (-150.0)) return image_np, label_np这里最关键的是SetSize的计算方式。如果目标 spacing 是 1mm×1mm×5mm而原图是 2mm×2mm×5mm那么新 size 应该是原来的spacing / target_spacing倍。注意 z 方向的层数可能因为四舍五入差一层这时 label 和 image 的 size 必须完全一致否则后续 np 数组不齐训练时 loss 维度直接报错。窗宽窗位 -150 到 250 是腹部软组织的常规设置能覆盖肝、脾、肾的灰度范围。如果某些病例里有明显的胆囊高密度或胰腺钙化这个范围也基本能包住。做完 clip 后需要做归一化到 0~1很多 TransUnet 开源实现会默认对输入做这种归一化如果训练效果差先检查这里是不是重复归一化。3.3 标签映射与类别处理数据集的标注值可能是连续的 1~9也可能是非连续的比如没有胆囊的病例里 label 4 完全不存在。很多损失函数实现会对torch.eye(num_classes)[label]做 one-hot如果类别值不连续生成的 one-hot 张量维度会错乱。所以进入训练前我会先把标签重映射成 0~N-1。# 把 Synapse 的 1~9 映射成 0~8方便 one-hot label_map { 1: 0, 2: 1, 3: 2, 4: 3, 5: 4, 6: 5, 7: 6, 8: 7, 9: 8 } remapped np.zeros_like(label_np, dtypenp.uint8) for src, dst in label_map.items(): remapped[label_np src] dst如果你研究的重点是“肾脏”而不是“左肾/右肾”可以把左右肾合并合并后类别数变成 8。注意合并时不能简单把 2 和 3 都改成同一个值因为其他类别可能是 4~9改成同一个值后需要再做一次压缩。最省事的办法是先合并再重新从小到大编号。3.4 生成文件列表与固定划分我见过不少人每次训练前用random.shuffle重新划分数据集导致两个 epoch 之间训练样本完全不同模型收敛不稳定也无法和其他方法公平对比。正确做法是固定随机种子并生成文件列表后续每次实验都读同一个train.txt和val.txt。import os import random random.seed(42) cases os.listdir(data/imagesTr) cases.sort() # 先排序再 shuffle避免跨系统目录序不同 random.shuffle(cases) train_cases cases[:18] val_cases cases[18:] with open(train.txt, w) as f: for c in train_cases: f.write(fdata/imagesTr/{c} data/labelsTr/{c}\n) with open(val.txt, w) as f: for c in val_cases: f.write(fdata/imagesTr/{c} data/labelsTr/{c}\n)有个容易踩的地方是路径中既有图像也有标签如果标签文件名有_label后缀文件列表里要写清楚顺序要固定为image_path label_path。之后在 Dataset 里直接用两列路径读图不要靠字符串拼接推理路径否则遇到文件名不规则的病例就会错位。4. 从零训练基于 PyTorch 的 TransUnet 腹部多脏器分割实战4.1 最小训练代码结构这里的训练代码按常见开源实现来写重点不是复现完整模型而是把数据加载、模型初始化和训练循环串起来。我习惯用out_channels9还是10取决于你背景是否单独成类Synapse 数据集通常把背景算作 0所以类别数等于 9。import torch import torch.nn as nn from torch.utils.data import DataLoader from transunet_model import TransUnet # 换成你手头的实现 model TransUnet( img_dim224, in_channels1, out_channels9, # 背景 8 个器官或按你重映射后的类别数 hidden_dim768, n_layers12, n_heads12 ).cuda() criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) for epoch in range(300): model.train() total_loss 0.0 for images, labels in train_loader: images, labels images.cuda(), labels.cuda() logits model(images) loss criterion(logits, labels.long()) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item()这段代码能跑通但请注意CrossEntropyLoss在这里只是个基础版本它会把 90% 的梯度投向背景。训练到第 50 个 epoch 时可能平均 Dice 只有 0.5 上下所以多数人会加一个 Dice loss 来拉平衡。out_channels的数值要和预处理里的重映射对应不要模型里写 10标签却只有 0~8。4.2 必调参数batch size、学习率、损失权重、类别权重TransUnet 的训练敏感点比纯 U-Net 多。我习惯用 AdamW 而不是 Adam学习率初始 1e-4配合余弦退火到 300 个 epoch。batch size 在 11GB 显存下单卡用 8 比较合适输入再大的话就需要自动混合精度。损失函数上用 Dice loss 和交叉熵各 0.5Dice loss 里忽略背景。参数推荐值说明输入尺寸224×224再大 token 数翻倍显存压力大batch size811GB 可用16 需要 AMP优化器AdamW对 Transformer 更稳初始学习率1e-43e-4 以上容易发散学习率策略CosineAnnealing300 epoch 后降到接近 0损失函数Dice CE 各 0.5Dice 部分忽略背景训练轮数300验证集指标到 200 epoch 后基本平缓类别权重我一般不给 CE 设置权重而是在 Dice loss 中直接按类别计算让每个器官的贡献相等。这样肝脏和胰腺在损失函数里的地位是相同的不会因为肝脏面积大而主导梯度。如果你手动给胰腺调高权重也要小心增幅过大否则模型会在胰腺上过拟合整体平均 Dice 反而下降。4.3 训练结果怎么看Dice、HD95 和预测可视化训练结束后我习惯对每个器官单独计算 Dice再算一个平均 Dice。只看整体 Dice 会被背景稀释一个模型背景 Dice 0.99 但前景全错整体可能还是 0.9 以上。计算时应该排除背景或者只对target 0的位置求交并比。import numpy as np def dice_score(pred, target, cls): # 第 cls 类的 Dicecls0 表示背景但统计前景时不看 0 pred_mask (pred cls) targ_mask (target cls) inter np.sum(pred_mask targ_mask) total np.sum(pred_mask) np.sum(targ_mask) return (2.0 * inter 1e-6) / (total 1e-6)在 Synapse 常见 18/12 划分下一套能用的结果大致是这样的肝脏 0.94、脾脏 0.91、左肾 0.92、右肾 0.91、胰腺 0.61、胆囊 0.78、食管 0.75、主动脉 0.87平均 0.84。胰腺和胆囊明显低于肝脏这是正常现象。如果你看到胰腺只有 0.2先别急着调模型回看预处理里胰腺区域是否被窗宽窗位裁掉了或者 label 和 image 是否在 z 方向错了一两层。4.4 保存最优模型与恢复训练我遇到过训练到 150 个 epoch 时验证集 Dice 到了 0.83但最后保存的 last epoch 反而只有 0.79。因为验证集很小最后一个 epoch 可能碰上一例难例就被拉低。所以我在训练过程中用best_dice保存一个最优 checkpoint而不是直接保存最后一步。best_dice 0.0 for epoch in range(300): # ... 训练循环 ... val_dice validate(model, val_loader) if val_dice best_dice: best_dice val_dice torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_dice: best_dice, }, best_transunet.pth)保存整个 dict 而不是只保存模型参数是为了以后恢复训练时还要恢复优化器的学习率和动量。如果你想做更长时间的实验甚至可以每 10 个 epoch 存一次避免断电后从头开始。这个 checkpoint 就是你的后悔药。5. TransUnet 腹部分割的常见问题排查Dice 虚高、显存溢出和标签错位5.1 显存不足batch size4 还是 OOM现象训练很快中断CUDA out of memory反复出现。原因输入图虽然显示是 224×224但有些病例在重采样后 z 轴层数很多一个 batch 里切片总数暴涨或者你的img_dim设置过大导致 Transformer 的 token sequence 太长。解决先把输入尺寸固定到 224×224再开启混合精度。混合精度可以用torch.cuda.amp包裹 forward 和 loss 计算对 Dice loss 几乎无副作用。还不行就把 batch size 降到 4并用梯度累积模拟 batch size8。5.2 训练很快收敛但预测结果全黑现象训练 loss 下降很快验证集 Dice 也有 0.8 以上但看可视化预测图时整个 mask 都是黑的或者只有零星几个点。原因背景像素占绝对多数模型学会了“什么都不预测”也能让 Dice 很高。如果你只看整体 Dice 不看前景 Dice这个问题会被掩盖。解决损失函数改成 Dice CE并对 Dice loss 部分把背景通道剔除。评估时也只统计目标类别大于 0 的 Dice不要统计背景。另一个快速验证手段是打印验证集中的一例预测看胰腺或肝脏区域是否有输出。5.3 标签错位训练 loss 正常验证指标却乱跳现象训练 loss 稳步下降但验证时甲病例的胰腺预测和乙病例的肝脏重叠或者某个器官形状严重扭曲。原因最典型的是重采样后图像和标签的 z 轴层数不一致。图像用线性插值后尺寸变了标签用最近邻插值后尺寸也变了但如果SetSize两次不一致数组切片就会错位。解决在预处理函数最后加一个断言assert image_np.shape label_np.shape并可视化一个 case 的中间层把 image 和 label 叠加后看一眼。如果错位检查两个 ResampleFilter 的SetSize是否都基于同一个resampled_img.GetSize()。5.4 数据增强过度左右翻转让模型彻底学乱现象加了随机水平翻转后验证集 Dice 不但没涨反而比不过不加增强的版本。原因腹部 CT 中左右肾、脾脏和肝脏的空间关系是解剖学固定的左右翻转相当于把器官位置互换Transformer 学到的空间先验被破坏了。解决不要用水平翻转和 90° 旋转。只保留随机亮度、小幅度缩放和平移。如果要做更复杂的数据增强建议用医学图像常用的弹性形变但形变幅度要小避免把一个器官拉成另一器官的形状。5.5 验证集 Dice 波动大不知道该取哪个 epoch现象第 200 个 epoch 验证集 Dice 是 0.85第 210 个 epoch 就掉到 0.78下一个 epoch 又回到 0.84。原因验证集只有 12 例单次 epoch 的噪声本身就大。另一个原因是验证时只做一次前向没有对数据做任何平滑或 TTA遇到一例难样本就会剧烈波动。解决用 4 中提到的best_dice做 checkpoint 选择最后用这个最优模型重新跑一遍验证集并记录每个器官的 Dice。不要以最后一个 epoch 的输出来写报告也不要因为某个 epoch 突然升高就提前结束训练。要判断是否收敛看近 20 个 epoch 的平均 Dice 而不是单点值。6. 让模型结果更稳验证集策略、TTA 与器官类别权重如果你不满足于跑出一个能看的数值接下来要做的是让结果在多次实验中稳定。第一步是改验证策略我在 Synapse 这类小数据集上会用 5 折交叉验证而不是只做一次 18/12 划分。每次用 24 例训练、6 例验证共训练 5 次最后报告平均 Dice 和标准差。这样能有效避免某一次划分里恰好把难例全分到验证集导致模型评分偏低。对于时间紧张的场景至少也要换两个固定随机种子跑两遍看看结果差异大不大。TTA 这块我不建议做水平翻转但可以做多尺度推理。具体做法是同一张切片分别用 224×224 和 192×192 输入尺寸送进模型把 softmax 概率图平均后再取 argmax。小尺度推理能看到更大范围大尺度能保留更多细节两者融合后对小器官的稳定性有帮助。代价是推理时间翻倍适合在验证集上观察不适合在线上服务里每张图都做。类别权重方面我建议直接在每个类别上单独算 Dice 并求平均这等于变相给小器官更高的权重。如果你一定要在 CE loss 里加权重可以按1 / log(freq 1)设置避免小器官权重过大导致过拟合。下表是我常用的比例可参考器官相对体积建议权重倍数肝脏大1.0脾脏中1.2肾脏中1.2胆囊小1.8胰腺小2.0最后分享一个我的习惯每 10 个 epoch 保存一次验证集的预测切片图挑出失败的那几张单独看而不是只看 Dice 曲线。肝脏看起来圆润、胰腺却碎成几段这就是明显的感受野不足或预处理把胰腺裁掉了胆囊颜色和肠壁混在一起就检查窗宽窗位。指标会骗人图像不会尤其是胰腺这类小器官跑完看一眼预测图比调十个参数都管用。希望帮到你。本文还有配套的精品资源点击获取