TransUNet：当Transformer遇上UNet会擦出怎样的火花？（医学图像分割新思路）

本文链接：https://blog.csdn.net/javajedi1/article/details/148014346

文章目录

前言（为什么你还在用传统UNet？）

各位搞CV的小伙伴们注意了（敲黑板）！今天要给大家介绍的这个网络结构，绝对能让你的分割模型性能原地起飞！！！传统UNet在医学图像处理领域称霸了这么多年，是时候来点新玩法了——TransUNet这个将Transformer和UNet巧妙融合的网络，在胰腺分割任务上直接把Dice系数干到了87.6%！（比原版UNet高了整整9.8个百分点）

一、TransUNet结构大拆解（Transformer的正确打开姿势）

1.1 网络整体架构（这个设计太妙了！）

整个网络就像个三明治结构（见下方示意图），最底层是CNN特征提取器，中间夹着Transformer层，最上层是UNet风格的解码器。这种设计既保留了CNN的局部特征捕捉能力，又通过Transformer获得了全局上下文信息。

重点来了（必考知识点）：

输入图像先被切成16x16的patch（跟ViT的处理方式类似）
使用ResNet-50的前4个stage作为编码器（别问为什么不用VGG，问就是残差连接真香！）
关键创新点：在CNN特征图上叠加位置编码后送入Transformer（这个操作让模型既懂空间位置又懂语义信息）

1.2 Transformer编码器详解（不是简单的堆叠！）

这里的Transformer层可不是随便堆的（新手最容易踩的坑）！！！作者采用了12层的Transformer blocks，每层包含：

Multi-Head Attention（8个注意力头）
MLP扩展比为4的前馈网络
层归一化（LayerNorm）和残差连接

注意（超级重要）：在医学图像中，病灶区域往往只占很小部分，所以这里的注意力机制要重点关注局部细节和全局位置的关系！

二、手把手实现TransUNet（PyTorch实战篇）

2.1 环境准备（别在版本问题上翻车！）

# 必备的三件套
import torch
import torch.nn as nn
from einops import rearrange  # 张量操作神器！

# 版本建议（血泪教训）：
# PyTorch 1.7+ / torchvision 0.8+ / timm 0.4.5+

2.2 核心代码实现（逐行解析）

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