Swin-Unet这篇论文我盯了很久MICCAI 2022上的工作把Swin Transformer和U-Net结构缝在一起做医学图像分割。当时看到标题的第一反应是“又是Transformer缝合怪”但仔细读完代码和实验后发现它在处理不规则目标、小器官分割上的确有两把刷子跟纯CNN的U-Net系列相比属于思路完全不同的另一条路。这篇文章我把整个模型的来龙去脉、网络结构、关键参数、训练心得和踩坑记录一次性理清楚适合已经跑过基础分割模型、想换Transformer架构试水的朋友也适合刚入门但想直接上手Swin-Unet的读者。1. 整体设计思路拆解为什么是Swin Transformer U-Net1.1 医学图像分割真正需要什么传统U-Net的编码器-解码器结构统治医学分割快十年了靠的是CNN的平移等变性和局部感受野。但CNN的天花板很明显卷积核尺寸有限想捕捉长距离依赖就得堆深度、加空洞卷积结果要么参数量爆炸要么梯度传递困难。尤其是在器官边界模糊、目标形状高度不规则比如肝脏肿瘤、脑部小病灶的场景下局部特征往往不够用需要更大范围的上下文信息。Swin Transformer的shifted window机制把自注意力限制在局部窗口内同时通过窗口移位实现跨窗口信息交互。这里有两个特别吸引人的点一是计算复杂度从ViT的二次方降到了线性输入图像可以做到更大尺寸二是它天然具备层级多尺度特征跟U-Net的下采样-上采样结构能一一对应。Swin-Unet做的就是把这两套东西拼起来编码器用Swin Transformer块替换CNN卷积块解码器用patch expanding层替换转置卷积中间保留skip connection做特征融合。1.2 选型时对比过的其他方案我在读这篇论文之前还对比过TransUNet和nnFormer。 TransUNet的做法是先用CNN提取特征再在最后一层加一个Transformer编码器本质上还是CNN为主、Transformer为辅长距离依赖只在最高层捕获。nnFormer则直接在3D体数据上用Transformer做编码解码但它的attention机制复杂很多训练不稳定对显存的要求也高。Swin-Unet相对更“纯粹”一些整个编码器全是Transformer块没有CNN参与而且因为用的是shifted window attention输入尺寸的限制没有ViT那么大。我拿2D腹部CT切片试过输入512x512的图Swin-Unet能跑得动TransUNet在同尺寸下显存占用就明显偏高。另外Swin-Unet还能借用Swin Transformer在ImageNet上预训练好的权重做初始化对医学小数据集来说这个优势很关键。1.3 这个设计解决了什么问题长距离依赖建模能力窗口注意力可以逐步扩大感受野高层窗口能看到低层窗口的信息等于在空间维度上做了全局建模。多尺度层级特征patch merging逐层下采样天然形成尺度金字塔跟U-Net的skip connection设计天然匹配。训练稳定性相比ViT直接做全局attention局部窗口的inductive bias更温和不容易在医学图像这种数据量小的任务上崩。实际跑下来的感觉是在病灶面积小、边界模糊的场景比如胰腺、肾脏肿瘤Swin-Unet的分割结果比U-Net要锐利不少不容易出现“多切一块”或者“少切一块”的偏差。2. 网络结构核心细节与关键参数2.1 整体架构一览Swin-Unet的基本流程是输入图像先经过patch embedding切成patch序列然后依次经过4个stage的编码器每个stage由patch merging和若干个Swin Transformer block组成瓶颈层保持分辨率不变解码器用patch expanding逐级恢复分辨率每级解码器输出跟对应编码器的skip connection拼接。以输入224x224x3的RGB医学图像为例Stage 1patch size4得到56x56x96的特征图经过2个Swin Transformer blockStage 2patch merging后分辨率减半到28x28通道翻倍到192经过2个blockStage 3继续下采样到14x14x384经过6个blockStage 4下采样到7x7x768经过2个block。解码器这边对称patch expanding每次把分辨率恢复为两倍通道减半skip connection把编码器的同分辨率特征concat进来补细节。2.2 Shifted Window Attention到底在做什么这是Swin-Unet最核心的部分也是跟普通U-Net最大的区别。标准的多头自注意力是每个像素跟所有其他像素做attention计算量是序列长度的平方。shifted window attention先把特征图切成固定大小的窗口论文默认window size7只在窗口内部做自注意力这样计算量就变成了窗口数量乘以窗口内像素数的平方大幅度下降。但只做窗口内注意力有个问题窗口和窗口之间没有信息交流。Swin的解决办法是下一层把窗口向右下方向平移shift让上一层的窗口边界变成这一层的窗口内部相当于通过“层间错位”来实现跨窗口信息传递。具体实现上作者用了cyclic shift的技巧把移出边界的部分绕回到对面省去了padding的计算浪费。实际写代码时有个容易出错的细节cyclic shift后需要给移动过来的区域加attention mask否则原本不相邻的像素会被错误地当成邻居。Swin-Unet的官方代码里有一大段生成mask的逻辑刚接触时很容易看懵我建议直接复用官方实现不要自己重写。2.3 Patch Merging和Patch Expanding的实现细节Patch merging的作用是下采样做法是把2x2邻域的四个patch在通道维度上concat再经过一层线性映射把通道数调整为新分辨率下想要的维度。比如输入是56x56x96经过patch merging变成28x28x192空间分辨率减半通道数翻倍。Patch expanding正好反过来有两种常见实现方式一种是先做线性映射把通道数翻倍然后用rearrange操作把通道维拆成空间维实现2倍上采样另一种是直接使用nn.Upsample加上卷积。原论文用的是第一种它在前向传播时能更好地保留空间信息不会像转置卷积那样产生棋盘效应。2.4 关键参数速查表参数默认值说明输入尺寸224x224可按数据集调整需能被patch size整除Patch size4初始patch切分尺寸Window size7窗口注意力窗口大小Embedding dim96C第一层特征通道数Stage层数[2, 2, 6, 2]各stage的Transformer block数量Head数量[3, 6, 12, 24]多头注意力的head数MLP ratio4MLP隐藏层维度是embedding dim的4倍Dropout0.1默认dropout比例QKV biasTrue注意力计算中是否加入bias这个配置是Swin-Tiny级别的模型参数量大约27M-30M左右比标准U-Net大但比ViT-Base小很多单卡训练压力并不大。3. 实操过程与训练配置从数据准备到模型收敛3.1 环境与代码选择如果只是想复现论文跑个效果直接起Swin-Unet官方仓库的代码PyTorch实现。我的建议是不要自己从零搭Swin Transformer block因为shifted window的mask逻辑细节太多出错了很不好排查。官方代码虽然风格有点学术化但胜在经过了大量验证稳定性没问题直接用就可以了。环境依赖方面我的测试配置是Python 3.8 PyTorch 1.10CUDA 11.2单张NVIDIA RTX 309024GB显存Swin-Unet的显存占用大头在self-attention和MLP层24GB显存跑batch size8、224x224输入没有压力。如果显存只有12GB建议把batch size降到4或者把输入裁剪到192x192。3.2 数据准备与预处理医学图像分割的数据集格式通常比较乱有DICOM的、有NIfTI的、也有直接导出成PNG的。Swin-Unet官方代码默认输入是通用的二维图像所以第一步就是把原始数据统一处理成png或者npy格式的切片。我自己做腹部CT分割时用了一个预处理流程把CT的HU值窗口化到[-200, 200]过滤掉无关组织归一化到[0,1]按固定分辨率resize到224x224保存为2D切片对应的mask做同样的resize注意用最近邻插值不能用线性插值否则标签边界会糊掉。数据量方面如果只有几十例标注建议做online数据增强随机旋转、水平翻转、随机裁剪、亮度扰动。Swin-Unet对数据量还是比较敏感的增强策略直接决定了最终效果的上限。3.3 损失函数与优化器选择官方代码用的损失函数是CrossEntropyLoss DiceLoss的组合这是一个很稳妥的选择。CE负责像素级别的分类正确性Dice负责区域重叠度两者加权相加能兼顾边界准确度和区域完整性。权重我一般设CE1.0Dice1.0如果类别不平衡严重比如病灶只占全图的1%可以考虑把Dice权重提高到2.0。优化器我用的是AdamW初始学习率5e-4weight decay设为1e-2。这里有一个经验Transformer类的模型对学习率其实比CNN敏感lr设太大会直接训飞设太小又收敛很慢。我一般先用一个小数据集跑5个epoch测试观察loss曲线是否平滑下降如果波动太剧烈就降一半学习率。训练轮次上医学小数据集通常100-200个epoch就能收敛加上warmup策略会更快。我习惯前10个epoch用线性warmup从1e-6升到设定学习率后面用余弦退火衰减到1e-6。3.4 混合精度训练与显存优化Swin-Unet在训练时用自动混合精度AMP是必须的能省将近一半显存而且因为注意力计算对精度不敏感几乎不会掉点。PyTorch里直接用torch.cuda.amp就行。需要注意的是Layernorm和softmax这两个操作建议保留FP32精度AMP的autocast会自动处理但如果你手动混精度要格外小心这两个地方。另外DiceLoss在FP16下容易出现梯度不稳定我在代码里把损失计算放在FP32下做而不是FP16。显存实在不够的情况下还可以用gradient checkpointing重算激活值显存能再降30%但训练时间会增加把window size从7降到5显存下降很明显但感受野会变小减少batch size并增加梯度累积步数效果等价但BN层不会受影响——不过Swin-Unet用的是LayerNorm所以梯度累积完全没问题。4. 训练过程中遇到的问题与排查实录4.1 损失函数不下降或收敛极慢这个问题出现的频率非常高大部分情况是学习率设置不合理导致的。Transformer结构的模型跟CNN不一样它对学习率的容忍度低很多在我自己跑的过程中AdamW配合5e-4的初始lr是一个相对安全的起点。如果把学习率降到了1e-4还是损失曲线平得跟直线一样那就要检查数据预处理了——最常见的问题是mask和原图对不上或者归一化把图像信息压倒在一个非常窄的范围内。还有一个容易忽略的点Swin-Unet的初始权重如果用的是ImageNet预训练第一轮epoch的loss通常不会太低因为医学图像和自然图像的特征分布差异很大。不要因为前几个epoch的loss没有快速下降就直接放弃给它10个epoch的时间如果还是没有明显下降再考虑调整学习率或者数据增强策略。4.2 显存爆掉OOM显存爆掉基本集中在两处attention计算和MLP层。我的排查步骤是先把batch size降为1如果还爆说明是模型本身的显存占用就太大需要缩小输入分辨率或window sizebatch size1能跑起来的话逐步增加batch size找到当前显存能承受的上限检查是否有不必要的中间变量被保存比如在推理时开了梯度或者把整张图padding到很大的尺寸。Window size的影响比想象中大从7降到5attention部分的显存占用能减少约一半。但要注意window size改变之后输入尺寸也要跟着调整保证能被window size整除否则代码会在padding或unpatch时报错。4.3 分割结果边界粗糙或出现空洞Transformer模型容易出现分割结果比较“碎”的情况尤其是小目标。我的经验是出现这个问题先不要急着调模型结构优先检查损失函数。纯CrossEntropy在小目标上几乎注定会失败因为背景像素占了绝大多数模型只需要把所有像素预测为背景就能得到很低的loss。加了DiceLoss会好很多但还是可能出现“猜测性”小空洞。另一个实用技巧是在解码器最后一层输出之后加一个CRF后处理或者简单的形态学闭运算。这个不是论文里写的但实践效果不错能把细碎的假阳性小区域清掉让边界更干净。如果追求端到端也可以在训练时用soft dice配合边界距离map效果更好但不一定值得额外折腾。4.4 训练一个多星期没进展久训不收敛除了学习率问题还有一个很常见的原因预训练权重的加载不对。Swin-Unet的权重是Swin Transformer在ImageNet上预训练好的但医学图像是单通道灰度图输入通道数跟ImageNet的3通道不一样。处理方式我试过两种把灰度图复制成3通道这样可以直接加载预训练权重论文里默认是这个做法修改第一层卷积的输入通道为1只加载后面层的权重。实测下来第一种方式明显更稳虽然增加了输入冗余但预训练带来的特征提取能力能完整迁移过来。第二种方式在第一层就截断了预训练的信息流效果会打折。这个经验我在多个数据集上都验证过。5. Swin-Unet的实际表现分析适合什么场景不适合什么5.1 跟U-Net相比的优势场景把Swin-Unet和U-Net在同一个数据集上对比我最有体感的一个场景是有多个同类目标且互相靠近的情况。比如肾脏肿瘤分割有时候左右肾的肿瘤在同一个切片上很近U-Net很容易把两个病灶连成一片Swin-Unet因为attention机制的全局建模能力能更好地分辨出这是两个独立目标而不是中间有个狭长连接。另外在肝脏血管分割这种长条状、蜿蜒结构的任务上Swin-Unet比U-Net表现好。CNN因为局部感受野的限制对连续细长结构的感知是断裂的而Swin Transformer的窗口注意力通过层级整合能更好地保持这种长程连续性分割出来的血管不会断成一截一截的。5.2 不适合的场景和局限Swin-Unet并不是无所不能。最基本的限制就是它本身是2D网络处理3D体数据时只能逐切片推理会丢失切片间的空间上下文。如果数据本身就是3D CT或MRI直接用nnUNet这种3D CNN架构往往是更实用、效果更好的选择Swin-Unet在3D数据上的优势没那么大。另外Swin-Unet的训练成本确实比U-Net高。在同样数据量下U-Net可能只需要50个epoch就能达到不错的效果Swin-Unet可能要100个epoch以上。如果项目周期很紧、没有充足的标注数据U-Net依然是更稳妥的选择。5.3 关于改进方向的思考如果想让Swin-Unet在特定数据集上再提一个档次我试过几个方向效果都不错在skip connection上加一个简单的注意力门控让模型学到底层特征和深层特征哪些部分更应该融合把最后一个stage的window size增大等于在最高层做全局注意力提升大目标的分割一致性训练时用多尺度输入随机把输入resize到0.8倍或1.2倍增强对目标尺度变化的鲁棒性。这些方法不用改核心网络结构只是在现有代码上做小幅调整但带来的收益通常比调损失函数更明显。6. 个人实际使用中的一些心得Swin-Unet给我最大的启发是医学图像分割的下一步不见得是设计更深的CNN而是怎么把序列建模的优势用起来。局部窗口注意力天然就是为“大图、小目标、精细边界”这类医学图像难题设计的虽然它也有训练成本高、数据需求大的弱点但在数据量足够、任务复杂度高的情况下确实能比传统U-Net系列多拿几个点。如果让我给一个建议新手第一次跑Swin-Unet不用在结构上纠结太多先把官方代码跑通用一个公开数据集比如Synapse多器官分割数据集复现论文效果。等你对shifted window、patch expanding这些操作有了直观感受再往自己的数据集上迁移会顺很多。预训练权重一定要用这会省掉大量调参时间。最后再分享一个小技巧如果你要在自己的数据集上微调Swin-Unet把图像读进来之后先看一下模型预测的可视化结果不要只看loss数值。很多问题比如mask错位、类别不均衡、增强过度单看loss是发现不了的可视化一眼就能看出问题在哪。这个习惯帮我排查掉了一半以上的“模型不work”问题。