各位做医学图像分割的朋友今天聊一个绕不开的模型——TransUNet。如果你手头正好有CT、MRI或者病理切片的分割需求又觉得纯CNN的U-Net虽然好但总是差点全局语义信息那这篇文章应该能给你一些启发。TransUNet是2021年发表在NeurIPS上的工作核心就一句话把Transformer当成编码器的增强组件嵌入到U-Net框架里用CNN提取底层空间特征用Transformer建模长距离依赖最后通过解码器逐步恢复分辨率。它在多器官分割、心脏分割等任务上的表现确实证明了这套混合架构的价值。这篇博文不打算只复述论文我会结合自己复现和改造TransUNet的实际经验把它的设计逻辑、实现细节、训练技巧和容易踩的坑都摊开来讲。不管你是刚入门的学生还是已经在做落地项目的工程师只要你用得上医学图像分割这篇文章应该都能给你一些实实在在的参考。1. 为什么偏偏是TransUNetU-Net的边界与Transformer的机会1.1 纯U-Net在医学分割里的瓶颈在深入TransUNet之前得先搞清楚它要解决的问题是什么。U-Net从2015年提出至今一直是医学图像分割的主力军原因很简单它在标注数据极其有限的情况下依然能靠数据增强和编码-解码结构学出不错的结果。但用过U-Net的人应该都有体会随着层数加深特征图在不断下采样的过程中会丢失不少空间细节。虽然跳跃连接能带来一些补偿但由于每个卷积核的感受野始终有限U-Net很难对相距很远的像素之间建立关系。举个例子在肝脏分割任务里肝脏区域的边缘通常和周围组织灰度对比并不强烈机器很容易把相邻的胃壁或肌肉组织误认成肝脏。如果网络只盯着局部纹理就非常容易产生这种边界混淆。想让模型真正理解“这是一个完整的肝脏”就必须让它对整张影像有一个全局级别的感知。传统的做法是把网络加深、加宽或者引入空洞卷积这虽然能扩大感受野但计算开销上去了很多小显存的卡根本跑不动。1.2 Transformer带来的可能性与它的尴尬Transformer在自然语言处理领域的成功让很多人看到了它在视觉任务上的潜力。它的自注意力机制能在一层之内让任意两个位置直接交互也就是说不管两个像素在空间上隔得多远模型都有能力学习它们之间的关联。这一点恰恰弥补了CNN的短板。但这里也有一个尴尬的问题Transformer在处理高分辨率图像时计算量是随序列长度的平方增长的。医学图像通常都是512x512、甚至1024x1024级别的单通道或三通道图如果你直接把每个像素当成一个token送进Transformer那计算量是任何一张显卡都承受不起的。所以ViT的做法是把图像切分成16x16的小块每一块当成一个token。但这样一来想要做到像素级精细分割Transformer对空间细节的恢复能力就不够了。1.3 TransUNet的设计思路不是二选一而是混合TransUNet的核心聪明之处在于它没有在CNN和Transformer之间做非此即彼的选择而是把两者组合在一起。具体来说先用一个CNN作为骨干网络提取特征得到一层分辨率适中、语义丰富的特征图然后把这张特征图切成一个个patch送入Transformer编码器去建模全局关系最后把经过Transformer处理的特征重新拼成二维特征图通过U-Net风格的解码器逐级恢复分辨率。这个设计的巧妙之处在于CNN部分负责搞定细粒度的空间结构Transformer部分负责搞定全局语义而解码器部分利用跳跃连接把底层细节和高层语义融合起来。我可以直接说这个组合在多个医学分割基准上的确比单纯的U-Net、单纯的TransUNet变体以及很多CNN与Transformer简单拼接的模型效果都要好。它并不是一个理论完美、但实操很难受的模型相反它对显存的要求和训练时间都处在可以接受的范围这也是它能被广泛复现和使用的原因。2. 深入网络骨架编码器、解码器与跳跃连接的设计考量2.1 CNN编码器用什么骨干网络决定了特征质量TransUNet整篇论文最容易被忽视、但实操中影响最大的部分其实是Encoder的CNN骨干网络选择。论文里实验了两种主流选择一种是经典的ResNet50另一种是轻量级的MobileNetV2。作者在不同实验里用了不同的backbone来验证他们的方法。我自己跑下来的感受是选择ResNet50作为骨干时模型整体的分割精度会比轻量级网络高一截但训练时间和显存占用都会明显增加。如果你在资源有限的场景下工作MobileNetV2版本的TransUNet仍然有不错的可用性尤其适合快速验证和部署。MobileNetV2靠深度可分离卷积大幅减少了参数量在差不多精度的情况下模型大小可能只有ResNet50版本的一半都不止。我的建议是如果你面向的是落地推理、边缘设备部署那优先选MobileNetV2版本如果你做的是学术实验、目标是在排行榜上刷精度那ResNet50是更稳妥的选择。2.2 Transformer编码器Patch嵌入和自注意力的细节Transformer编码器部分的输入来自CNN骨干网络输出的特征图。比如输入图像是224x224经过ResNet50多次下采样后我们拿到的是14x14的空间尺寸、通道数为512或768的特征图。接着要将这个特征图切分成一个个patch并用一个线性层做patch嵌入。14x14的特征图可以被切成1x1的patch也可以切成2x2的patch。如果patch大小越大得到的token数量就越少计算量越小但局部信息的粒度就越粗。论文里大量实验使用1x1的patch因为此时14x14等于196个token计算量相当小而且能在保持较高分辨率的同时让Transformer捕捉全局关系。你可以在自己的实验中按需调节这个参数但我个人的经验是医学图像的病灶结构通常比较紧凑1x1的patch输出效果往往最稳。Transformer层里还有两个关键的细节位置编码和类别token。位置编码用来告诉模型每个token在空间上的相对位置如果缺了这个信息Transformer就很难区分两张内容相同但位置不同的特征图。论文使用的是标准可学习位置编码训练时直接作为参数优化。而类别token是ViT里用来做图像分类的在TransUNet里因为任务不是分类这个token最终会被丢弃或掩码掉重点放在特征重组部分。2.3 解码器和跳跃连接空间分辨率的恢复路径经过Transformer编码器输出的特征序列需要重新reshape回二维特征图比如从196个token恢复成14x14的空间形状。然后进入解码器。TransUNet的解码器设计大体沿用U-Net的结构每一级先通过上采样操作放大分辨率再将来自编码器同层级的特征拼接进来融合后经过卷积层精细调整。如此反复直到恢复成和原始输入相同的分辨率。之所以要拼接编码器特征是因为高层特征虽然有丰富语义但空间位置是很粗糙的底层特征虽然语义弱但边缘和纹理信息保留得很好。两者结合网络既能知道“这块区域是个器官”又能定位“这个器官的边界在哪个像素”。这个设计的另一个好处是缓解了Transformer带来的局部细节损失。Transformer在建模全局关系时并不会天然保留像素级的边界信息跳跃连接等于在恢复阶段把这一层信息从底层直接引回来避免解码器只靠高层特征而丢失边缘。3. 从公式到代码损失函数、训练配置与复现实验3.1 输入尺寸与预处理并非越大越好先把最实际的训练配置讲清楚。医学图像分割的输入尺寸需要根据GPU显存情况去权衡。我最初尝试的时候直接用512x512的分辨率输入期望保留尽可能多的细节结果一个batch只能放两张图训练速度慢得让人崩溃而且小目标组织分割的精度并没有明显比224x224高很多。后来我做了实验对比在Synapse多器官分割数据集上224x224输入配合resize和随机裁剪分割效果和384x384输入非常接近但训练时间是后者的三分之一不到。原因在于多器官分割任务中各个器官的空间分布和边界在较低分辨率下其实已经能够被较好估计更关键的是语义信息而不是微观纹理。如果你的显存只有12GB左右建议优先使用224x224并使用一些在线数据增强来弥补信息损失如果你有24GB以上的显存可以直接上384x384甚至更高。预处理方面医学图像的窗宽窗位调整至关重要。以CT影像为例不同组织的CT值范围差异很大。一般建议将图像裁剪到感兴趣的窗宽范围内比如肝脏分割常在[-150, 250]之间再把像素归一化到[0, 1]或均值为0、方差为1。没有这一步模型会花费很多不必要的参数去区分不同患者之间的扫描差异。3.2 损失函数选择交叉熵、Dice Loss还是混合损失分类任务里大家习惯用交叉熵但医学分割任务有一个非常典型的问题前景和背景的像素数量极度不平衡。比如一个512x512的CT切片里肝脏可能只占几百个人像素背景却有几十万个。这时如果只用交叉熵模型很容易把所有像素都预测为背景来换取很低的loss。常见的解决方案是Dice Loss它直接优化分割结果和标注之间的重合度。它的公式很简单等于2乘以预测和标注的交集除以上两者的并集。Dice Loss对类别不平衡问题有很强的容忍性因为不管前景多小只要没预测准loss就会很高。但我实际用下来纯Dice Loss在训练初期非常不稳定因为梯度变化太剧烈导致网络收敛困难。更稳的做法是把Dice Loss和交叉熵组合起来。我常用的是Dice Loss 加权交叉熵权重分别设为0.5和0.5。这样既保持了交叉熵在像素级别上的平滑梯度又能利用Dice Loss加强对小目标的约束。在Synapse数据集的实验中这种组合比单独使用Dice Loss提高了约1.5%的Dice系数。3.3 训练策略从预训练权重开始的收益TransUNet的Transformer部分建议从预训练权重初始化。这个预训练通常是在ImageNet上用图像分类任务训练的ViT或DeiT权重。有些人可能会想医学图像和自然图像差异这么大直接用预训练权重会不会反而不好我的实测结果是即使考虑了这个差异从预训练权重初始化的模型在分割精度上仍然要明显优于随机初始化并且收敛速度快得多。原因也好理解Transformer的底层特征学习的是通用的视觉结构比如边缘、颜色、纹理组合这些特征在自然图像和医学图像之间是相通的。用预训练权重相当于让模型站在了巨人的肩膀上只花少量数据去适应医学图像的分布即可。建议在训练前先冻结骨干网络和Transformer编码器的前几层只改动解码器部分观察loss变化待整个模型有了一个基础拟合能力后再全部解冻这样能有效防止训练初期的大幅震荡。3.4 实验复现一个可以用作baseline的训练流程示例我自己复现时使用的训练框架是PyTorch配合两个GPUbatch size设为16初始学习率设成了1e-4采用AdamW优化器设置weight decay为1e-4。整个流程可以参考下面的伪代码结构model TransUNet( img_size224, in_channels1, num_classesnum_classes, embed_dim512, depth12, num_heads8, patch_size1, backboneresnet50 ) criterion CombinedLoss(ce_weight0.5, dice_weight0.5) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max200) for epoch in range(200): for images, labels in train_loader: images, labels images.cuda(), labels.cuda() outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() # 关于验证集上的Dice和IoU指标统计 eval_model(model, val_loader)训练过程中我建议每5个epoch保存一次checkpoint同时保留一个专门的best模型。很多情况下模型在后期会出现过拟合表现在验证集Dice不再上升但训练集loss还在下降。如果遇到这种情况果断回退到验证指标最好的那个权重或者提前终止训练。Synapse数据集上我这个配置大约在120个epoch左右就能达到一个稳定的分割效果。4. 从零到一实现时的硬核细节与性能优化4.1 训练内存优化混合精度与梯度累积医学图像分割模型训练内存占用高这是不少人劝退的原因。以TransUNet为例输入224x224、batch size为16的时候一张24GB的显卡可能会被挤到只剩几个G。如果batch size只能设为2或4那么模型训练的效果会明显变差因为Batch Normalization层的统计信息在很小的batch里会非常不稳定。一个直接的解决办法是使用梯度累积。简单说就是把若干个batch的梯度累加起来达到目标batch size对应的梯度后再执行一次参数更新。这样我们虽然每次只喂少量图像进模型但梯度的累积效果跟完整batch一样。另一个手段是开启PyTorch自带的混合精度训练也就是AMP。开启AMP后模型在前向和反向传播时自动将部分计算切换到FP16半精度显存占用能降一半左右训练速度还能提升不少。我在启用AMP和梯度累积后原来需要24GB显存才能跑起来的配置在12GB显卡上也顺利完成了训练。4.2 数据加载效率别让IO成为训练瓶颈很多人在模型上花了很多心思却忽略了数据加载这一环。尤其是在医学图像场景每次读取的都是NIfTI格式的三维数据或者整张全切片IO开销相当大。我最初用普通DataLoader时数据加载时间居然占了总训练时间的40%。需要针对性地做一些优化。首先是实现一个高效的数据加载类在__getitem__函数里做切片和预处理时尽量用numpy的切片操作而不是循环逐像素处理。其次是开启DataLoader的num_workers参数通常设置为CPU核心数的一半数据加载就能显著提速。另外如果内存足够可以一次把整个训练集读入内存在内存里完成切片和数据增强这是最省时的方案。对于超大切片建议预先做patches的切割把每张训练图像切成若干个固定大小的patch保存下来训练时直接读取patch文件。4.3 推理阶段加速减少冗余计算的思路训练结束后推理阶段的效率也很关键。TransUNet中Transformer编码器部分是计算量最大的模块之一。如果你要在多张图像上做推理可以尝试把Transformer编码器在训练时学习的权重固定不动只对图像做一次特征提取和一次全局建模。对于同一组参数不需要在每张测试图上重新计算前向过程中的某些临时变量PyTorch的torch.no_grad()模式可以节省大量不必要的显存和计算时间。另外一个更实际的加速策略是模型剪枝和蒸馏。可以把TransUNet当成教师模型训练一个结构更简单的学生模型来模拟教师模型的输出。这种方法在医学图像场景已经有了一些成功案例学生模型的速度能提升数倍同时精度损失控制在可接受范围内。不过这块内容比较深需要分数据去实验适合已经对TransUNet本身有充分理解的读者。5. 常见问题与排查技巧实录5.1 训练不收敛先检查这五件事如果你用TransUNet训练时发现loss几乎不下降或者Dice系数一直在零点几徘徊先别急着调模型结构按顺序检查下面五件事第一检查数据归一化是否合理。医学图像如果直接用原始像素值送入模型数值范围可能非常大导致梯度爆炸。CT图像尤其要注意窗宽窗位处理。第二检查标签编码。有些医学数据集的标注文件用的是255表示目标区域、0表示背景如果直接用交叉熵计算模型很难收敛正确的做法是把标签改成类别索引比如0、1、2。第三检查损失函数是否写对。Dice Loss的公式里分子分母如果计算维度不对很容易出现范围内浮点误差导致的NaN。第四检查学习率是否合适。Transformer对学习率比较敏感1e-4是个相对稳妥的起点如果你用了更大的学习率建议先降到3e-5再试试。最后检查类别权重是否设置正确。如果你对多器官分割任务按器官类别设置权重权重差距过大也会导致训练不稳定。5.2 分割结果出现空洞或边缘毛刺这类问题相信很多做分割的人都遇到过。模型输出的结果有时候会呈现不规则的杂点或者在器官内部出现被误分割成背景的小洞。这通常是因为模型在解码器的最后一层上特征图分辨率不够细腻加上损失函数对边界像素的约束不够严格。我的处理经验是在训练阶段增加边界感知的损失项。具体来说可以对标注图像做一次边缘提取把边缘像素的权重调得更高让模型更关注这些困难位置。还可以在最后输出层后增加一个条件随机场后处理能够有效整理分割边缘。不过CRF的推理速度偏慢如果是在线服务场景需要权衡。最直接的办法是调整patch大小和输入分辨率往往能把边缘质量提升一个档次。5.3 多个器官互相粘连边界区分困难Synapse这类多器官数据集上肝脏和脾脏、胃和胰腺之间的距离非常近灰度分布甚至很接近模型经常把它们连在一起。前期我跑的时候器官粘连问题特别突出测评的Dice指标倒还可以但看了一眼可视化结果发现两个器官完全糊成一片。解决办法之一是在标签处理上利用器官之间的相对位置信息。Transformer的一个优势就是能捕获这种远距离关系所以如果你发现模型还是粘连很有可能是Transformer层数不够深或者patch尺寸太大导致分割很粗糙。可以尝试加深深度比如从12层编码器变成16层并把patch从2x2改成1x1。我在实际实验中发现将patch设为1x1同时加深两层Synapse数据集上的脏器分离效果就有明显改善。6. 适用场景与扩展方向哪些任务适合用它哪些未必6.1 适合TransUNet的任务特点我总结了一下TransUNet特别适合那些既需要精细边界、又需要全局语义的任务。多器官分割、心脏结构分割、肿瘤区域分割、视网膜血管分割这些任务要么关注的目标尺度变化很大要么目标分布不太规则需要全局上下文辅助判断。TransUNet在这些场景下通常能比U-Net高出几个百分点的Dice得分。具体来说在心脏MRI分割任务中心肌的形状相对稳定但心腔的边缘和周围组织对比度低这时Transformer的全局感知能帮助减少误判。在肝脏和肝肿瘤分割任务中肿瘤的尺寸、形状变化极大纯CNN容易把零散小结节漏掉TransUNet的全局建模则能捕捉到肿瘤和非肿瘤组织的微妙差异。6.2 不一定适合的任务你该怎么判断TransUNet也不是万能的对于那种图像尺寸巨大、目标非常小、且目标之间没有明显空间关联的任务它的优势不一定明显但计算开销却很大。比如血管分割血管形态细长且在整个图像中分布稀疏Transformer虽然能覆盖全局但特征过粗时反而容易忽略细小分支。另一个不太适合的是实时推理场景。因为Transformer编码器的计算量毕竟比纯卷积大不少如果对单张图像的推理时间要求在毫秒级TransUNet带来的提升可能不值得牺牲的延迟。对这些任务可以考虑轻量化的变体比如把TransUNet编码器替换成更小的Tiny版本或者采用知识蒸馏的方式只学习主流器官的信息、降低解码器的复杂度。6.3 衍生方向TransUNet之后它还留下了哪些思路TransUNet本身只是一个开始它启发了大量后续工作这些工作基本都是在它的框架上做了局部创新。有的把CNN骨干替换成了更强大的ConvNeXt或Swin Transformer有的把Transformer编码器换成了Swin Transformer因为Swin使用窗口注意力机制计算效率和局部建模能力更强也有的引入了多尺度的思想通过不同层级的特征图分别计算Transformer attention再融合结果。如果你想在这个方向继续深入建议先吃透TransUNet的代码实现再去理解它的衍生模型会节省很多周折。我个人在实际操作中的体会是做医学图像分割拼模型结构是最简单的一步真正花时间的是数据处理、损失函数设计和调参。TransUNet提供了很好的骨架但最终效果的差异往往取决于你怎么处理数据、怎么设计训练流程、怎么解决落地场景里的各种脏问题。如果准备跑这个模型可以先拿Synapse或ACDC数据集跑通整个流程再迁移到自己的数据上这个路径相对平缓能让人少走很多弯路。