TransUNet：当Transformer遇上医学图像分割（实战指南）

本文链接：https://blog.csdn.net/progalchemist/article/details/148025398

文章目录

一、从U-Net到TransUNet的进化之路

“U-Net？那个对称的编码器-解码器结构？” 没错！这个2015年诞生的经典网络（作者Olaf Ronneberger绝对是个天才！）至今仍是医学图像分割的baseline。但（敲黑板）随着Transformer在CV领域大杀四方，2021年CVPR的最佳论文候选——TransUNet横空出世！

传统U-Net的痛点在哪？举个栗子🌰：当处理复杂器官边界（比如肝脏的不规则形状）时，CNN的局部感受野容易漏掉长程依赖。而Transformer的自注意力机制（Self-Attention）刚好能捕捉全局上下文信息！

二、TransUNet结构全解析（附网络示意图）

2.1 整体架构（灵魂画手上线）

[输入图像] → [CNN特征提取] → [Transformer编码] → [CNN解码上采样] → [输出分割图]

这个混合架构才是精髓！先用CNN提取局部特征（像医生先看局部病变），再用Transformer分析全局关系（类似专家会诊综合判断）

2.2 三大核心模块

Patch嵌入层（关键第一步！）
- 把CNN提取的特征图拆分成16x16的patch
- 代码实现技巧：nn.Unfold(kernel_size=16, stride=16)
Transformer编码器（注意力机制核心）
- 多头注意力层数不是越多越好！论文里用了12层，但实际项目6-8层足够
- 位置编码千万别忘！医学图像的空间关系至关重要
跳跃连接改进版（U-Net的灵魂升级）
- 传统U-Net直接拼接低阶特征
- TransUNet增加了通道注意力模块（SE Block）自动学习特征重要性

三、PyTorch实战代码（手把手教学）

3.1 环境准备

# 必备库安装（建议新建虚拟环境）
pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
pip install einops  # 张量操作神器！

3.2 Transformer编码器实现

class TransformerBlock(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4.):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, int(dim*mlp_ratio)),
            nn.GELU(),  # 比ReLU更适合Transformer
            nn.Linear(int(dim*mlp_ratio), dim)
        )
    
    def forward(self, x):
        # 残差连接必备！
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.mlp(self.norm2(x))
        return x

3.3 完整模型搭建要点

class TransUNet(nn.Module):
    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()
        # CNN编码器（ResNet34为例）
        self.encoder = resnet34(pretrained=True)
        # Patch嵌入
        self.patch_embed = nn.Conv2d(512, 768, kernel_size=16, stride=16)
        # Transformer序列处理
        self.transformer = nn.Sequential(*[TransformerBlock(768, 12) for _ in range(12)])
        # 解码器上采样
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(768, 512, 3, stride=2),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            # ... 逐层上采样
        )

四、训练技巧大公开（血泪经验）

4.1 数据增强策略

空间变换：随机旋转（-15°~15°）、弹性变形（模拟器官蠕动）
颜色扰动：只在HSV空间调整明度（保持医学图像颜色特性！）
混合增强：CutMix比CutOut更有效！

4.2 损失函数选择

推荐Dice Loss + BCE联合使用：

class DiceBCELoss(nn.Module):
    def __init__(self, smooth=1e-6):
        super().__init__()
        self.smooth = smooth

    def forward(self, pred, target):
        pred = torch.sigmoid(pred)
        intersection = (pred * target).sum()
        dice = (2.*intersection + self.smooth) / (pred.sum() + target.sum() + self.smooth)
        bce = F.binary_cross_entropy(pred, target)
        return bce + (1 - dice)  # 反向组合

4.3 学习率设置（关键！）

使用Warmup+Cosine衰减：

from torch.optim.lr_scheduler import CosineAnnealingLR

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scheduler = CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)
# 前5个epoch线性warmup
for epoch in range(5):
    optimizer.param_groups[0]['lr'] = 1e-4 * (epoch+1)/5