简介:面向医学图像分割与深度学习入门进阶人群,这份实战资料围绕基于TransUNet的腹部多脏器分割任务,涵盖背景、肝脏、左/右肾、脾脏五类结构,提供完整可运行的代码、数据集与训练结果。代码采用PyTorch实现,包括训练、评估、推理三类脚本,并配备详细注释与README说明,方便替换自定义数据后快速上手。包体包含1031个文件,以986个PNG图像(原图与掩膜)、18个Python脚本、2个PTH权重文件为主,另有训练日志、txt配置与说明文档,整体约200.83MB。训练环节使用AdamW优化器、余弦退火学习率衰减与交叉熵损失,自动输出loss/IOU曲线、学习率衰减曲线、数据集可视化及最优权重。100个epoch训练后,测试集像素准确率达到0.986,平均IOU为0.779,可直观感受模型效果与调参结果。已有536人学习下载,适合希望理论结合代码、快速开展腹部多脏器分割实验的研究者与开发者。
1. TransUnet 腹部多脏器分割:为什么这个组合值得你动手跑一遍
一份腹部 CT 进来,要同时把肝脏、脾脏、左肾、右肾、胰腺、胆囊、胃这些器官逐个画出来,这是典型的语义分割任务,而且是多类别、类别严重不平衡的那种。用纯 U-Net 跑,小器官经常漏检;用纯 ViT 跑,边界又糊得没法看。TransUnet 把两者串起来——ViT 做编码器抓全局上下文,CNN 解码器恢复空间细节,正好卡在医学影像分割的痛点上。这篇文章要拆的就是这样一份完整资源:代码、数据集、训练好的权重都齐,从数据预处理讲到训练调参,再把推理验证的流程走一遍。适合两类人:一是想做新架构复现的算法工程师,二是医学影像方向需要拿公开数据集出结果的学生。照着跑通一次,你对 TransUnet 的理解会比只看论文深得多。
2. 数据准备:从原始 NIfTI 到 TransUnet 能吃的输入
2.1 数据集选择与目录结构
腹部多脏器分割最常用的公开数据集是 Synapse 多脏器分割数据集,CT 序列以 NIfTI 格式存储,每例包含全腹部扫描和对应的 8 个器官标注:主动脉、胆囊、左肾、右肾、肝脏、胰腺、脾脏、胃。训练集 18 例、测试集 12 例,单个病例的切片数量大约在 85 到 200 张之间,原始分辨率是 512x512,层厚 5mm 左右。
拿到这份资源后,先把目录结构理清楚,我的习惯是统一改成下面这样:
project/ ├── data/ │ ├── train/ │ │ ├── img/ │ │ └── label/ │ └── test/ │ ├── img/ │ └── label/ ├── code/ │ ├── dataset.py │ ├── model.py │ ├── train.py │ └── predict.py └── results/ ├── weights/ └── logs/逻辑说明:把原始 NIfTI 按训练集和测试集拆开放,图片和标签分层存放,后面写 Dataset 类时路径拼接非常省事。results 目录单独放权重和日志,训练过程中每轮的模型备份和评估指标都落在这里,方便回溯。
参数说明:数据集路径建议用相对路径,代码里用os.path.join拼接,避免换机器后硬编码路径失效。权重目录按训练日期命名,比如results/weights/20250612_transunet/,我习惯在模型文件名后缀加上 epoch 和验证 Dice,一眼能看出哪个权重最优。
2.2 预处理三件事:窗宽窗位、归一化与切片重采样
医学 CT 和自然图像最大的区别是灰度值范围。CT 值的单位是 Hounsfield Unit(HU),腹部软组织一般在 -100 到 300 之间,而原始数据范围可能到 -1024 到 3071。直接喂给网络,大部分灰度信息会被背景淹没。
常见的做法是先做窗宽窗位调整,把腹部窗口的 HU 范围映射到 0 到 1,再切成长方形的 2D 切片,最后统一 resize 到网络输入尺寸。下面这段是我在用的预处理代码:
import nibabel as nib import numpy as np import cv2 def preprocess_ct(nii_path, out_shape=(224, 224), window_level=40, window_width=400): img = nib.load(nii_path).get_fdata() # 窗宽窗位截断 lower = window_level - window_width / 2.0 # 40 - 200 = -160 upper = window_level + window_width / 2.0 # 40 + 200 = 240 img = np.clip(img, lower, upper) img = (img - lower) / (upper - lower) # 映射到 [0, 1] img = img.astype(np.float32) slices = [] for i in range(img.shape[2]): slice_2d = img[:, :, i] slice_resized = cv2.resize(slice_2d, out_shape, interpolation=cv2.INTER_LINEAR) slices.append(slice_resized) return np.stack(slices, axis=0) # [D, H, W]逻辑说明:先用np.clip把 HU 值截断在窗宽范围内,低于下界的赋值为下界、高于上界的赋值为上界,这样软组织之外的骨骼和空气不会干扰归一化。然后做 min-max 归一化,把窗口内的灰度映射到 0 到 1 之间,网络输入分布稳定。最后沿轴向逐切片截取并 resize。
参数说明:window_level=40, window_width=400是腹部软组织常用的窗口参数,覆盖了肝脏、脾脏、肾脏的灰度范围。如果你做的是骨分割或者肺部分割,窗口参数要换成骨窗(level 400、width 1800)或肺窗(level -600、width 1500),这是很多新手直接套用自然图像归一化导致分割效果差的主要原因。
2.3 标签编码与数据增强
标签文件是整张 NIfTI,像素值为 0 到 8 的整数——0 是背景,1 到 8 对应 8 个器官。这里有个关键点:多类别分割不需要把标签做成 one-hot 存盘,直接保留整数标签,训练时在损失函数里做 one-hot 转换,可以省不少磁盘空间和 I/O 时间。
但切片尺寸很小(224x224),要防止过拟合,得加数据增强。我用的增强策略是同步空间变换,确保图像和标签做完全相同的变换:
import random def augment(image, label, p=0.5): if random.random() < p: # 水平翻转:图像和标签同步翻转 image = np.flip(image, axis=1).copy() label = np.flip(label, axis=1).copy() if random.random() < p * 0.6: # 随机旋转 ±10 度,用相同旋转矩阵 angle = random.uniform(-10, 10) image = rotate_image(image, angle) label = rotate_image(label, angle, is_label=True) return image, label逻辑说明:翻转和旋转是医学图像分割里最安全的空间增强,不会改变器官的相对位置关系。标签用最近邻插值而图像用双线性插值,因为最近邻能保持标签的整数类别值不变,双线性插值会在类别边界产生小数,导致后续损失计算出错——这是很隐蔽的坑。
参数说明:旋转角度 ±10 度对腹部器官来说是合理范围,超过 15 度会产生大量背景区域,反而干扰训练。翻转概率 p 设 0.5,让数据分布均匀。颜色增强如亮度、对比度扰动在 CT 上不需要做,因为 CT 值本身已经通过窗宽窗位校准过,随意扰动会破坏灰度物理意义。
3. 模型拆解:CNN 编码器与 ViT 是怎么拼接的
3.1 为什么说 TransUnet 是「U-Net 的骨架、Transformer 的大脑」
纯 U-Net 的问题在于卷积操作的局部性——每一层卷积的感受野有限,虽然下采样能扩大感知范围,但连续卷积在编码高维语义时对长距离依赖的建模能力始终受限。纯 ViT 则相反,self-attention 天然建模全局关系,但 patch embedding 把图像切碎后会丢失精细的空间结构,直接上采样恢复的分割边界毛糙。
TransUnet 的思路是分两段走:先用 CNN 做浅层特征提取,保留高频细节;再把这些特征图展平成 token 序列送入 ViT 做全局建模;最后通过 U-Net 风格的解码器和跳跃连接恢复分辨率。这样既有了 ViT 的全局感受野,又保留了 U-Net 的精确边界恢复能力。
3.2 Patch Embedding 与 ViT 编码器的代码视角
看这份资源的 model.py,核心结构是「ResNet50 前几层 + ViT 主干」。图像先通过一个 stem 卷积把分辨率降四倍,然后切成固定大小的 patch 送入 Transformer。关键代码长这样:
import torch import torch.nn as nn class PatchEmbedding(nn.Module): def __init__(self, in_channels=3, embed_dim=768, patch_size=16): super().__init__() self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) def forward(self, x): # x: [B, 3, 256, 256] x = self.proj(x) # [B, 768, 16, 16] B, C, H, W = x.shape x = x.flatten(2).transpose(1, 2) # [B, 256, 768] return x, (H, W)逻辑说明:patch embedding 的本质是patch_size x patch_size的无重叠卷积。以 256x256 输入、patch_size 16 为例,投影后变成 16x16 个 token,每个 token 是 768 维向量。这里的 768 是 ViT-Base 的默认隐藏维度,注意它不是超参随便设的——embed_dim 太大显存翻倍,太小全局建模能力下降。
参数说明:这份资源默认输入尺寸是 224x224,patch_size 是 16,token 序列长度 196。显存不够时优先改输入尺寸而不是 patch_size——patch_size 改成 32 会让空间细节直接减半,分割结果明显变粗。如果你有 8GB 显存,把输入缩到 192x192 或 160x160 比降 batch size 效果更稳。
3.3 解码器与跳跃连接:全局特征如何回到像素空间
ViT 输出的 token 序列是 [B, 196, 768],要重新变成 2D 特征图才能和 U-Net 解码器对接。这里有个关键细节:很多人直接 reshape 回 [B, 768, 14, 14],但 TransUnet 会在 reshape 之前加一个卷积层把维度通道对齐到解码器需要的通道数。
class DecoderBlock(nn.Module): def __init__(self, in_ch, skip_ch, out_ch): super().__init__() self.up = nn.ConvTranspose2d(in_ch, out_ch, kernel_size=2, stride=2) self.fuse = nn.Conv2d(out_ch + skip_ch, out_ch, kernel_size=3, padding=1) def forward(self, x, skip): x = self.up(x) # 上采样 2 倍 x = torch.cat([x, skip], dim=1) # 跳跃连接拼接 x = self.fuse(x) return x逻辑说明:每个解码块先反卷积上采样两倍,再把 CNN 编码器对应层的特征图拼接进来。这个 skip connection 是 U-Net 架构的灵魂,让解码器在恢复空间细节时可以直接参考浅层的高分辨率特征,不用完全凭空生成。
参数说明:通道数设计遵循 U-Net 的经典「减半」策略。如果编码器输出 512 通道,解码器第一层降到 256,拼接浅层 256 后融合成 256,依次类推,最后输出层用 1x1 卷积把通道数映射到类别数 9(背景 + 8 器官)。注意最后一层不要加 softmax——把 logits 直接交给损失函数,数值稳定性更好。
4. 训练配置与调参:从损失函数到收敛判断
4.1 损失函数:Dice Loss 与 Cross-Entropy 的加权组合
腹部多脏器分割最典型的痛点是类别不平衡。肝脏占腹部切面面积可能超过 20%,胆囊可能只有不到 1%。这时候用纯 CrossEntropy Loss,模型倾向于把所有像素预测成背景和肝,胆囊、胰腺这种小器官直接消失。
这份资源用的损失函数是 Dice Loss 和 CrossEntropy 的加权组合,代码核心如下:
class DiceCE(torch.nn.Module): def __init__(self, num_classes=9, ce_weight=0.5, dice_weight=0.5): super().__init__() self.num_classes = num_classes self.ce_weight = ce_weight self.dice_weight = dice_weight def forward(self, logits, targets): # logits: [B, 9, H, W], targets: [B, H, W] ce = torch.nn.functional.cross_entropy(logits, targets) probs = torch.softmax(logits, dim=1) # [B, 9, H, W] targets_onehot = torch.nn.functional.one_hot( targets, num_classes=self.num_classes ).permute(0, 3, 1, 2).float() # 逐类别计算 Dice smooth = 1.0 intersection = (probs * targets_onehot).sum(dim=(0, 2, 3)) union = probs.sum(dim=(0, 2, 3)) + targets_onehot.sum(dim=(0, 2, 3)) dice = (2.0 * intersection + smooth) / (union + smooth) dice_loss = 1.0 - dice.mean() return self.ce_weight * ce + self.dice_weight * dice_loss逻辑说明:Dice Loss 直接优化分割质量的核心指标,对类别不平衡天然鲁棒——它按类别计算重叠率再取平均,小器官的梯度不会被大器官淹没。CrossEntropy 提供逐像素的梯度,有助于模型快速收敛到大致区域。两部分的权重是可调的,类别越不平衡,dice_weight 应该越大。
参数说明:smooth=1.0是平滑项,防止分子分母同时为 0 导致数值异常。如果你发现某个类别完全没有被预测出来,可以把平滑项减小到 1e-5,让梯度信号更敏感。Dice 损失在训练早期数值很大是正常现象,不要因为 loss 值高就调小学习率。
4.2 训练参数推荐表
| 参数 | 推荐值 | 说明 |
|---|---|---|
| 输入尺寸 | 224x224 | 显存不足时先降到 192 或 160 |
| Patch Size | 16 | 不宜改成 32,空间细节损失严重 |
| Batch Size | 8 | 8GB 显存可用 4,配合梯度累积 |
| 优化器 | AdamW | 权重衰减 1e-4,比 Adam 泛化更好 |
| 初始学习率 | 1e-4 | 配合 Warmup,前 10 个 epoch 线性升温 |
| Scheduler | 余弦退火 | T_max=200,ETA_min=1e-6 |
| 训练轮数 | 200 | 验证 Dice 连续 20 轮不涨则早停 |
| 类别数 | 9 | 背景 + 8 个器官 |
4.3 训练命令与收敛判断
这份资源提供了完整的训练脚本,数据路径改好后直接启动即可:
python train.py \ --data_root ./data \ --exp_name transunet_synapse \ --img_size 224 \ --batch_size 8 \ --lr 1e-4 \ --epochs 200 \ --num_workers 4 \ --gpu 0逻辑说明:训练过程中,每个 epoch 结束后会在验证集上计算平均 Dice。收敛判断标准我习惯看两个信号:验证集 Dice 不再上升且波动小于 0.5%,同时训练 Dice 和验证 Dice 之间的差距小于 5%——说明模型既没欠拟合也没严重过拟合。
参数说明:num_workers设 4 到 8,数据读取瓶颈在磁盘 I/O,NIfTI 格式解压比较慢。如果数据放在机械硬盘上,建议先跑一次预处理把数据转成 npy 或 png 格式,训练速度能提升 3 倍以上。--gpu指定设备,多卡用户可以用--gpu 0,1配合 DataParallel。
5. 避坑指南:腹部多脏器分割的五个翻车点
5.1 显存不足(OOM)
现象:batch_size 设 16,训练第二个 epoch 直接报 CUDA out of memory。
原因:TransUnet 的 ViT 部分显存开销远大于纯 CNN。self-attention 的注意力矩阵是 O(n^2) 复杂度,输入 224x224、patch 16 时 token 数量 196,注意力矩阵本身占用的显存很高。加上 encoder 和 decoder 的通道数大,图存的梯度占用量比 U-Net 高一个量级。
解决:优先把输入尺寸从 224 降到 192 或 160,显存占用会按平方下降。其次把 batch_size 降到 4,开启梯度累积(accumulation_steps=2),用时间换空间。最后才考虑改 patch_size——它影响的是分割精度,改了之后要重新评估。
5.2 小器官完全没预测出来
现象:训练 100 轮后,肝脏、脾脏的 Dice 能到 0.85 以上,但胆囊的 Dice 一直是 0,预测图里完全没有胆囊区域。
原因:类别不平衡。胆囊在腹部 CT 里的像素占比经常不到 0.5%,CrossEntropy 的梯度信号被大器官淹没,Dice Loss 虽然对不平衡鲁棒,但如果初始预测就没有该类别,Dice 梯度接近 0,模型永远学不会。
解决:在损失函数基础上,对少样本器官的类别做梯度放大。具体做法是在计算 Dice 时给每个类别加权重,胆囊、胰腺这类器官权重设为 3 到 5,肝脏这种大器官权重设为 0.5。另一个有效手段是过采样——把包含胆囊的切片多喂几次,让模型在训练中更频繁地看到它。
5.3 Loss 持续下降但 Dice 不涨
现象:训练 loss 从 0.8 稳定降到 0.3,但验证集 Dice 卡在 0.72 不动了,连续 30 轮没有提升。
原因:这不是模型不收敛,而是数据对齐出了问题。常见嫌疑有两个:一是 NIfTI 的切片方向没统一,有的病例是轴向扫描,有的是冠状位重采样,切片切出来器官形状差异巨大;二是窗宽窗位没有逐病例校准,导致同一个器官在不同病例里的灰度分布不一致。
解决:预处理时打印每个病例的 shape 和 spacing 信息,确认轴向一致。窗宽窗位参数不要写死,按病例的肝脏区域统计灰度分布,用中位数和标准差动态计算窗口范围。这个坑排查最耗时,我的血泪经验是第一周跑通基线时先把所有病例的统计信息输出到一张表里,眼过一遍再训练。
5.4 推理时标签错位:预测的器官颜色全乱了
现象:训练和验证 Dice 都正常,但推理输出的预测图上,肝脏的位置显示成了脾脏的颜色。
原因:类别映射表对不上。Synapse 数据集的标签编号是 1 主动脉、2 胆囊、3 左肾、4 右肾、5 肝脏、6 胰腺、7 脾脏、8 胃,但训练脚本里可能用了不同的编号顺序。数据预处理阶段如果做了类别重映射,推理时忘记应用同样的映射,就会整体错位。
解决:把类别映射表单独提出来,训练和推理共用同一个配置文件。不要在两张脚本里各写一份映射字典——两份必然改一份漏一份。建议在预测输出后做一次计数校验:统计预测结果中各类别像素占比,如果背景占比超过 90% 或者某个器官像素数严重超出预期,立刻检查映射表。
5.5 数据增强导致标签漂移
现象:加了旋转增强后训练 loss 变大了 30%,而且验证集 Dice 波动剧烈。
原因:图像和标签的增强没有用同一个随机种子。random.random()在图像分支和标签分支各调用了一次,导致图像旋转了 10 度、标签旋转了 8 度,两个矩阵冲突——这是增强代码最常见也最难觉察的 bug,损失函数不会报错,但模型学到的永远是错位的对。
解决:图像和标签共用一次随机数的生成。我的做法是把增强函数改成接受同一个随机状态,或者用np.random生成一次旋转角度,图像和标签都用这个角度变换。写完增强函数后一定要做可视化验证:随机取 10 个切片,把 image 和 label 叠图画出来,肉眼检查边界是否对齐。
6. 推理与验证:Dice 与 HD95 还不够,切片叠图才是最终标准
推理流程本身很短:加载训练好的权重,对测试集逐切片跑前向,拿到 logits 后取 argmax 得到预测类别图,再计算指标。下面这段是评估脚本的核心:
def evaluate(model, dataloader, device): model.eval() dice_per_class = np.zeros(9) counts = np.zeros(9) with torch.no_grad(): for images, labels in dataloader: images = images.to(device) logits = model(images) preds = torch.argmax(logits, dim=1).cpu().numpy() labels = labels.numpy() for c in range(1, 9): mask_gt = (labels == c) mask_pred = (preds == c) denom = mask_gt.sum() + mask_pred.sum() if denom == 0: continue dice = 2.0 * (mask_gt & mask_pred).sum() / denom dice_per_class[c] += dice counts[c] += 1 dice_per_class = dice_per_class / np.maximum(counts, 1) return dice_per_class逻辑说明:逐类别计算 Dice 而不是算全局 Dice,因为全局 Dice 会被肝脏这种大器官撑高,掩盖小器官分割失败的问题。每个类别的 Dice 单独累加再取平均,才能真实反映模型在所有器官上的表现。
参数说明:counts数组防止某些类别在某个 batch 里完全没出现导致除零。完整的评估还应该算 HD95(95% 豪斯多夫距离),衡量边界误差,代码里用scipy.ndimage.distance_transform_edt计算距离变换后取 95 分位数即可。HD95 对前列腺、胆囊这类边界不规则的器官比 Dice 更敏感。
评估完指标后,我强烈建议做一次切片叠图检查:把原始切片、预测掩膜、真实标签三张图叠在一起,用不同颜色标注不同器官,逐切片滑过去用肉眼检查。指标反映的是整体数值,叠图反映的是细节质量——比如肝脏边缘有没有锯齿状伪影、脾脏有没有把周围脂肪误判进去、胃的腔内边界是否完整。这些信息是任何单一指标都表达不出来的。
这个验证习惯是我从一次翻车经历里总结出来的。上次训练结果显示平均 Dice 到了 0.87,我以为万事大吉,结果切片叠图发现脾脏区域预测出了一个细长的假阳性条带,恰好对着胃的左外侧缘,指标被大器官的贡献掩盖了。从那以后我每次训练完都强制走一遍逐切片叠图检查,先看全貌再看异常区域,最后才整理指标上报。希望帮到你——别让 0.87 的平均值骗了你的眼睛。
本文还有配套的精品资源,点击获取