简介：TransUnet复现完整代码，面向医学图像分割开发者与相关研究者，帮助理解Transformer全局注意力机制与U-Net卷积结构相结合的分割模型。资源提供一整套可运行实现，覆盖数据预处理、模型构建、训练、测试与结果可视化，便于掌握原理后迁移到自有项目中。压缩包共44个文件，约751.19MB，主体为14个Python脚本，包含训练入口、训练器、辅助函数、维度转换及标签可视化等模块；另有文本说明、Markdown文档、预训练权重和XML配置文件，可直接搭建并调整参数。目前已有6303人学习下载，适合具备一定深度学习基础、希望复现医学图像分割模型并深入理解其实现细节的开发者。借助完整代码与说明，读者可以快速跑通TransUnet流程，并对照注释与配置进行二次开发，提升对Transformer与U-Net融合设计的整体认知。

1. TransUnet复现：不是简单堆一个Transformer，而是把U-Net和ViT的互补性真正用起来

医学图像分割的从业者对U-Net都不陌生，但遇到小目标、边缘模糊、低对比度这些场景时，纯CNN的感受野限制会非常明显。TransUnet的思路是保留U-Net的编码-解码骨架，在底层引入Transformer分支，用ViT来建模全局依赖关系。这份复现代码不是那种“能跑通就行”的玩具项目，它把混合架构的每个细节都拆开讲清楚了——包括CTA模块里的特征重映射、四层跳连接如何对齐尺寸、混合损失函数为什么比单用CrossEntropy稳定。适合正在做医学影像、遥感分割或者任何需要长距离依赖建模的开发者，新手拿来入门ViT分割也合适，但前提是你得能跟上。

2. 前置理论：TransUnet的架构拆解，CNN与Transformer到底在哪一层融合

2.1 混合编码器的设计逻辑，为什么不是把ViT整个替换掉

纯ViT做分割时会把输入切块成patch序列，计算量可控但对局部细节不敏感；纯CNN在深层虽然能看到较大区域，但靠堆叠卷积的感受野增长是线性的。TransUnet的做法是从ResNet-50的第三和第四层分别取输出特征图，把第四层的特征图切块后送入Transformer分支，每层输出的序列长度按patch大小和特征图分辨率计算。

我把核心的编码器构建逻辑提取成了下面这段代码，项目中对应的实现文件说明了整个流程：

# 编码器：CNN特征提取 + Transformer全局建模 class TransEncode(nn.Module): def __init__(self, img_size=256, in_chans=3, embed_dim=768, depth=12, num_heads=12): super().__init__() # 使用ResNet-50作为CNN分支，截取layer2和layer3的输出 resnet = models.resnet50(pretrained=True) self.layer1 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool) self.layer2 = resnet.layer1 self.layer3 = resnet.layer2 self.layer4 = resnet.layer3 # 把layer4输出的特征图展平成patch序列 self.patch_embed = PatchEmbedding(in_channels=1024, patch_size=16, embed_dim=embed_dim) # Transformer编码器，depth表示堆叠的blocks数量 self.transformer = TransformerEncoder(depth=depth, embed_dim=embed_dim, num_heads=num_heads) def forward(self, x): x1 = self.layer1(x) x2 = self.layer2(x1) x3 = self.layer3(x2) x4 = self.layer4(x3) # 将CNN特征转成序列送入Transformer seq = self.patch_embed(x4) seq = self.transformer(seq) return x2, x3, seq

这段代码最重要的设计决策是：CNN分支保留layer2和layer3的输出不送进Transformer，只把layer4的特征图交给ViT。这样浅层的空间细节不会被全局建模稀释。embed_dim=768是ViT-Base的默认配置，深度12层意味着Transformer分支有12个自注意力block。

2.2 PatchEmbedding和位置编码的参数定义，序列化过程中最容易被忽略的坑

从特征图到序列这一步，很多人直接照搬ViT官方的patch切法，但在TransUnet里输入不是原始图像，而是已经被CNN下采样过的特征图，尺寸和通道数都变了。项目里的实现是先把1024通道的特征图用1x1卷积压缩到embed_dim，再做展平，而不是像ViT那样直接用16x16卷积核切原图。

class PatchEmbedding(nn.Module): def __init__(self, in_channels=1024, patch_size=16, embed_dim=768): super().__init__() self.patch_size = patch_size self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) def forward(self, x): # x: [B, 1024, H, W]，经过16x16卷积后变成 [B, embed_dim, H/16, W/16] x = self.proj(x) B, C, H, W = x.shape x = x.flatten(2).transpose(1, 2) # [B, H*W, C] return x

patch_size=16在这里基于的是最后一张特征图的尺寸，而不是原图尺寸。如果原图是256x256，经过四层下采样后是16x16，卷积后变成1x1，序列长度就是1，没有任何意义。项目里处理方式是调整下采样次数或者改变patch大小，让序列长度在16x16到32x32之间比较合理。位置编码是nn.Parameter初始化的可学习参数，大小为[1, num_patches, embed_dim]，如果你的输入分辨率跟预训练模型不一致，这个参数会被随机初始化，效果会打折扣。

3. 解码器与跳连接：从768维序列恢复到256x256分割图，特征对齐是核心

3.1 CTA模块的作用，为什么需要两层卷积而不是直接上采样

Transformer分支输出的序列要跟CNN浅层特征融合，直接reshape回特征图再拼接到解码器里会出问题——通道数不匹配、空间分辨率不一致、语义层级差异大。项目里的CTA（Cross-Transformer Attention）模块是解决这三层错配的关键组件，它先把Transformer输出的特征重塑成特征图，再用两层3x3卷积调整通道数，让特征图和解码器当前层的通道保持一致。

# CTA模块：把Transformer序列转回特征图并与CNN分支特征融合 class CTA(nn.Module): def __init__(self, embed_dim=768, conv_channels=512): super().__init__() self.conv1 = nn.Conv2d(embed_dim, conv_channels, kernel_size=3, padding=1) self.bn1 = nn.BatchNorm2d(conv_channels) self.conv2 = nn.Conv2d(conv_channels, conv_channels, kernel_size=3, padding=1) self.bn2 = nn.BatchNorm2d(conv_channels) self.relu = nn.ReLU(inplace=True) def forward(self, seq, cnn_feat): # seq: [B, num_patches, embed_dim] -> reshape回特征图 B, N, C = seq.shape H = W = int(N ** 0.5) x = seq.transpose(1, 2).reshape(B, C, H, W) # [B, C, H, W] x = self.relu(self.bn1(self.conv1(x))) x = self.relu(self.bn2(self.conv2(x))) # 与CNN支路的浅层特征相加，而不是直接拼接 return x + cnn_feat

这里有个细节：CTA输出和CNN特征用的是加法融合，不是concat。加法融合在这个位置的合理性在于，Transformer分支和CNN分支在各自处理完特征后，需要把两种表征在同一个特征空间里合并，concat会让通道数翻倍，增加解码器参数量的同时并不一定带来更好的效果。

3.2 四层上采样结构，每一层的输入输出通道和特征图尺寸

解码器从最深层开始逐级恢复分辨率，最深层用的是512通道的CTA输出，依次向上通道减半、分辨率翻倍。项目中的解码器结构保持了U-Net的经典设计，每层包含两个3x3卷积，中间穿插ReLU和BatchNorm，上采样方式用的是双线性插值。

class Decoder(nn.Module): def __init__(self, base_channels=64, num_classes=2): super().__init__() # 四个解码层，通道数从512逐级降到64 self.up1 = DecoderLayer(in_channels=512, skip_channels=256, out_channels=256) self.up2 = DecoderLayer(in_channels=256, skip_channels=128, out_channels=128) self.up3 = DecoderLayer(in_channels=128, skip_channels=64, out_channels=64) self.up4 = DecoderLayer(in_channels=64, skip_channels=64, out_channels=64) self.final = nn.Conv2d(64, num_classes, kernel_size=1) def forward(self, cta_out, skip1, skip2, skip3): # cta_out来自Transformer分支，skip1-3来自CNN编码器 x = self.up1(cta_out, skip1) x = self.up2(x, skip2) x = self.up3(x, skip3) x = self.up4(x, None) return self.final(x)

DecoderLayer内部会先把上采样后的特征图和跳连接特征做concat，跨层特征对齐这一步最容易出问题。如果上采样尺寸和跳连接特征图尺寸差1个像素，concat会直接报错。项目里给出的经验是不要信任计算出来的尺寸，在forward里打印每一层feature map的shape，因为同样的卷积配置在不同输入尺寸下产生的特征图分辨率可能不一样。

4. 完整训练流程：混合损失函数、数据增强策略与评估指标

4.1 损失函数为什么选CrossEntropy加Dice Loss，权重配比的经验值

TransUnet在分割任务里容易出现前景背景像素极度不均衡的情况，肿瘤区域可能只占整张图的2%。单独用CrossEntropy会让模型倾向于把所有像素都预测为背景，单独用Dice Loss又有梯度不稳定的问题。项目的训练脚本用的是两者的加权和，Dice Loss负责处理类别不均衡，CrossEntropy负责稳定梯度和加速收敛。

# 混合损失函数：CrossEntropy + Dice Loss class MixedLoss(nn.Module): def __init__(self, ce_weight=1.0, dice_weight=1.0, num_classes=2): super().__init__() self.ce_weight = ce_weight self.dice_weight = dice_weight self.num_classes = num_classes def forward(self, logits, targets): ce = F.cross_entropy(logits, targets) # 计算每个类别的Dice系数，然后取平均 probs = F.softmax(logits, dim=1) dice = 0.0 for c in range(self.num_classes): prob_c = probs[:, c] target_c = (targets == c).float() intersection = (prob_c * target_c).sum() union = prob_c.sum() + target_c.sum() + 1e-8 dice += (2.0 * intersection) / union dice = 1.0 - dice / self.num_classes return self.ce_weight * ce + self.dice_weight * dice

ce_weight和dice_weight都设为1.0在大部分数据集上表现稳定。如果你的数据集特别不均衡，把dice_weight调高到2.0或3.0能更快看到效果，但训练前期loss可能会震荡，需要配合较小的学习率。

4.2 数据增强与归一化的参数设置，随机裁剪尺寸会影响模型精度

医学影像数据量通常不大，几百张图训练一个ViT规模的分割模型很容易过拟合。项目的数据加载部分做了随机旋转、水平翻转、随机亮度和对比度调整。归一化用的是ImageNet的均值和标准差，这里有个实际问题：如果你的输入图像是单通道灰度图，需要先复制成三通道再归一化，否则BN层的统计值会异常。

# 数据加载与增强配置 train_transform = transforms.Compose([ transforms.RandomResizedCrop(256, scale=(0.8, 1.0)), # 随机裁剪，尺度范围0.8~1.0 transforms.RandomHorizontalFlip(p=0.5), # 水平翻转概率 transforms.RandomRotation(degrees=15), # 随机旋转 transforms.ColorJitter(brightness=0.2, contrast=0.2), # 亮度和对比度抖动 transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

scale=(0.8, 1.0)这个参数很关键。如果裁剪尺度范围过小，模型每次看到的都是局部放大区域，分割边缘会很粗糙；范围过大又会让小目标区域被裁掉。0.8到1.0是保留上下文又不至于太多无关背景的经验区间。推理时不做RandomResizedCrop，固定resize到256x256就行。

4.3 训练监控指标：IoU和Dice的计算方式，跟损失函数里Dice的区别

训练脚本里同时计算IoU和Dice两个指标，用来判断模型是否真的在收敛。损失函数里的Dice是可微的近似版本，评估时的Dice是精确计算：按阈值把预测结果转成类别标签，再跟真实标签比对。

# 评估指标计算 def compute_iou_dice(pred, target, num_classes=2): pred = torch.argmax(pred, dim=1).flatten() target = target.flatten() ious = [] dices = [] for c in range(num_classes): pred_c = (pred == c) target_c = (target == c) intersection = (pred_c & target_c).sum().float() union = (pred_c | target_c).sum().float() iou = intersection / (union + 1e-8) dice = (2 * intersection) / (pred_c.sum() + target_c.sum() + 1e-8) ious.append(iou) dices.append(dice) return torch.mean(torch.tensor(ious)), torch.mean(torch.tensor(dices))

注意计算IoU时要加平滑项1e-8，避免某个类别在batch里完全没出现时除以0。这个评估函数的执行频率建议不低于每两个epoch一次，因为只靠loss曲线很难判断分割质量的真实变化。

5. 避坑指南：维度不匹配、显存溢出和效果不升反降的排查方案

5.1 输入尺寸改动后，位置编码和序列长度对不上

现象：把项目默认的256x256输入改成512x512后，前向传播报错，提示位置编码的维度不匹配。

原因：位置编码是定义在patch序列上的，patch数量由输入尺寸和patch_size共同决定。256x256输入搭配16x16 patch能得到256个patch序列，改到512x512后patch数变成1024，之前初始化的位置编码参数长度不够。

解决：要么保持输入尺寸不变，要么在修改输入尺寸后重新初始化位置编码。代码里定位这个问题的位置在模型初始化部分，搜索pos_embed相关的nn.Parameter定义，把第一维改成新序列长度后重新训练。

5.2 Transformer分支训练显存爆掉，batch size根本调小不下去

现象：batch size设成4直接OOM，换成2勉强能跑但GPU利用率很低。

原因：12层Transformer的显存占用跟序列长度的平方成正比，序列长度为256时每个头的注意力矩阵是256x256，4个batch同时计算就占掉大量显存。

解决：项目里的训练脚本支持梯度累积，把有效batch size保持在8，实际每次只计算2个batch的梯度，累积4步再更新参数。这样Transformer分支的显存瓶颈不会影响整体训练。

5.3 加了Transformer分支之后，效果反而不如纯U-Net

现象：同样的数据集，去掉Transformer分支的纯U-Net在验证集上Dice达到0.85，加上Transformer后反而降到0.78。

原因：最常见的情况是数据集本身病灶区域很小、边缘锐利，全局建模带来的上下文信息不足以弥补局部细节损失。还有一个容易被忽视的原因：Transformer分支的输入特征图分辨率过低，patch切分后序列长度太短，全局注意力学到的全是粗糙的语义信息。

解决：用项目配套的消融实验脚本先跑一遍，看看单独用CNN分支解码和单独用Transformer分支解码的效果差异。如果是数据集的病灶区域本来就小，可以试试把patch_size从16改到8，增大序列长度，让Transformer能感知更细粒度的特征。

6. 验证落地：用小数据集跑通全流程，用热力图检查预测置信度

项目里附带了一个小规模的医学影像分割数据集的样本，总共有几百张。拿到手先把train.py里的epoch数改成一个很小的值，比如3个epoch，先确认整个流程从数据加载到评估指标输出没有断点，再跑完整训练。第一次跑通之后，我习惯用类似下面的脚本检查模型对每张图的预测置信度分布：

# 检查预测置信度分布是否异常 def check_prediction_confidence(model, dataloader, device): model.eval() confidences = [] with torch.no_grad(): for images, _ in dataloader: images = images.to(device) logits = model(images) probs = F.softmax(logits, dim=1) max_prob, _ = torch.max(probs, dim=1) confidences.append(max_prob.mean().item()) avg_conf = np.mean(confidences) print(f"平均置信度: {avg_conf:.4f}") if avg_conf < 0.6: print("置信度过低，模型可能没有收敛") elif avg_conf > 0.98: print("置信度过高，可能存在过拟合风险")

这个检查脚本能帮你快速判断训练状态。平均置信度稳定在0.75到0.95之间是比较正常的分割结果区间。如果出现过拟合，优先检查BN层的running_mean和running_var在验证集上的表现，再看训练集和验证集的Dice差距有多大。

项目里的可视化脚本还会把预测结果叠加到原图上导出，这个习惯我保持了很久——只看指标数字容易忽略边缘缺损和孔洞这一类肉眼可见的问题。从那以后，我每次训练完分割模型都会强制走一遍可视化流程，确认边缘质量和临床可用性，而不是只盯着Dice分数。希望帮到你。

本文还有配套的精品资源，点击获取