Swin Transformer全面解析:窗口注意力与层级结构在视觉任务中的实践
Swin Transformer 是我这几年做视觉模型时印象最深刻的一篇工作没有之一。它在 ICCV 2021 拿了最佳论文这个头衔一点都不虚。简单来说它解决了一个非常实际问题Transformer 进入视觉领域之后图像这种高分辨率、强局部相关的数据到底应该怎么喂给注意力机制才能既高效又不丢细节。我用它做过分类、检测、分割还踩过不少坑这篇就把 Svin 的技术内核、复现要点和实际选型思路一次性讲清楚适合刚接触视觉 Transformer 的读者也适合已经跑过 ViT 但想在密集预测任务上提升效果的人。1. 为什么Swin Transformer能横扫视觉任务1.1 ViT的全局注意力在图像上栽过跟头在 Swin 之前ViT 已经证明了 Transformer 在视觉上可行但它有一个很尴尬的问题计算复杂度太高。ViT 把图像切成 patch比如 16x16 大小然后拉成一个序列直接做全局自注意力。224x224 的输入序列长度是 196勉强能跑。但一旦换成目标检测、语义分割这类任务输入分辨率动辄 800x800 甚至更高序列长度会变成几千个 token注意力矩阵的大小随序列长度平方增长显存和算力直接爆炸。很多人以为 ViT 在小数据集上打不过 ResNet 是模型能力不行其实更核心的原因是计算代价限制了它只能在小 patch 数量上做文章。全局注意力的问题在于它假设任意两个位置之间都需要直接建模关系但对于自然图像来说距离很远、语义无关的像素之间根本没有必要做完整交互。这种“无差别全局建模”既浪费预算又让模型很难学到图像特有的局部模式。ViT 后来之所以能通过大规模预训练拉回精度本质上是靠海量数据和超长训练时间硬生生补上了归纳偏置的缺失这在真实项目里代价太高了。1.2 移动窗口既不丢局部细节又省算力Swin 的核心突破是提出了移动窗口注意力。思路非常直接把特征图切成一个个不重叠的小窗口注意力只在每个窗口内部计算。比如窗口大小 7x7那每个窗口里有 49 个 token注意力计算量从整图 token 数的平方变成窗口数量乘以 49 的平方复杂度从二次降到了线性。这就好比公司开会与其每次全员一起开效率极低不如先按部门分组讨论每个部门内部充分沟通再找时间跨部门同步。图像里的局部纹理、边缘、小物体本质上都依赖局部上下文窗口注意力天然契合这种特性。但直接分窗会带来新问题窗口和窗口之间老死不相往来信息无法跨窗口传递感受野就被锁死了。Swin 的巧妙之处在于引入移位窗口也就是相邻层的窗口整体偏移。上一层的不重叠窗口到了下一层会交错重叠原本无法交流的区块就这样建立了联系。这相当于部门内讨论完之后重新分组再讨论一轮信息就扩散开了。这个设计看似简单但组合效果非常好既保持了局部建模的高效又通过层间交错实现了全局信息流通。1.3 层级结构才是检测分割真正需要的Swin 还有一个容易被忽略但极其重要的设计它像 CNN 一样构建了层级特征金字塔。具体做法是每经过一个 stage就通过 patch merging 把相邻 2x2 的 patch 合并成一个特征图分辨率减半通道数翻倍。整体流程从原始图像开始分辨率逐层从 H/4 降到 H/8、H/16、H/32和 ResNet 的 C2 到 C5 一模一样。这个特性对目标检测和语义分割来说几乎是救命稻草。因为这类任务天然需要多尺度信息小目标要在高分辨率特征图上检测大目标要在低分辨率特征图上识别。ViT 那种恒定分辨率的单尺度输出做分类还好做检测分割就很别扭要么额外加很多解码结构要么干脆效果很差。Swin 通过层级设计直接输出金字塔特征可以无缝衔接到 FPN、UperNet 这些成熟框架上这也是它能在下游任务里迅速取代 ResNet 的重要原因之一。2. 核心技术细节拆解计算量、位置编码与模型结构2.1 注意力复杂度从 O(N^2) 降到 O(N)把计算量摆出来看会更直观。标准全局注意力的 FLOPs 大概是 4hwC^2 2(hw)^2C其中 h、w 是特征图的高和宽C 是通道数。第一项是线性变换第二项是注意力矩阵的计算它随着 hw 的平方增长。窗口注意力把特征图分成很多 MxM 的小窗口计算量变成 4hwC^2 2M^2 hw C注意这里的 hw 前面没有平方了。举个例子输入 224x224patch size 4第一层特征图是 56x56C96。全局注意力的第二项大约是 2 x 3136 x 3136 x 96已经上亿了。窗口注意力按 M7 算窗口内 49 个 token一共有 64 个窗口第二项大约是 2 x 49 x 3136 x 96只有全局的六十四分之一。这就是 Swin 敢处理高分辨率输入的底气。实际训练时用 8 卡 V100 跑 Swin-T 的 ImageNet 预训练显存占用比同分辨率 ViT-B 小很多可以说这个设计是真正从工程压力里长出来的。2.2 相对位置偏置为什么这么关键Swin 的注意力公式是 Attention(Q,K,V)SoftMax(QK^T/sqrt(d)B)V其中 B 是相对位置偏置表。B 的维度是 (2M-1)x(2M-1)表示窗口内任意两个 token 之间在水平方向和垂直方向上的相对距离组合。你可以把它理解成一张查找表两个 token 的相对位置偏移决定加多少偏置而不是用绝对位置编码。我在实验里验证过去掉这个 B 之后ImageNet 分类精度会掉 1 到 2 个点检测任务的影响更大。原因也不难理解自然图像是有平移等变异质性的物体出现在左上角还是右下角不应该改变它的语义但绝对位置编码会强行引入坐标依赖。相对位置偏置让模型更关注两个点之间的“距离感”和“方向感”这对捕获边缘、纹理这类局部结构非常重要。而且这张表在不同窗口之间是共享的参数量很小不会给模型增加多少负担属于典型的四两拨千斤设计。2.3 Swin-T/S/B/L 的配置与设计逻辑Swin 家族里最常用的是 Swin-T、Swin-S、Swin-B、Swin-L 四个规格。以 Swin-T 为例输入 patch size 是 4embed_dim 是 96四个 stage 的 Transformer Block 数量分别是 2、2、6、2窗口大小固定为 7。每个 stage 开始前先做 patch merging把 2x2 的 patch 合并成一个通道数翻倍所以四个 stage 的输出通道分别是 96、192、384、768。Swin-B 则把 embed_dim 提高到 128每个 stage 的 block 数量也增加整体参数量大不少。值得注意的细节是同一 stage 内 W-MSA 和 SW-MSA 是交替出现的。W-MSA 是普通窗口注意力SW-MSA 是移位窗口注意力一个 block 做窗口内建模下一个 block 做跨窗口交互成对出现才能保证信息流通。如果你看官方代码会发现 SW-MSA 的实现用了循环位移和 attention mask 来模拟窗口偏移而不是真的把特征图重新切一次这个优化能显著减少张量搬运的耗时是工程实现里很聪明的一手。3. 实操落地从分类到检测分割的完整方案3.1 两行代码用上预训练 Swin如果你只是想快速验证 Swin 在自己数据上的效果我建议直接走 timm 这条路线别自己从头实现既省时间又不容易出错。代码很简单import timm model timm.create_model(swin_tiny_patch4_window7_224, pretrainedTrue)这样一行就会自动下载官方在 ImageNet-22K 上预训练好的权重输入尺寸默认 224x224。要注意的是timm 不同版本对 Swin 的权重加载方式有细微差异比如相对位置偏置表的初始化升级库之后最好先跑一次单 batch 前向确保输出形状正常再继续训练。另外如果要用更高分辨率比如 384x384可以选swin_tiny_patch4_window7_384这个模型的 position embedding 和相对位置偏置都已经做过相应插值比在 224 模型上手动改输入尺寸更稳。但要注意切换输入分辨率后同样要检查相对位置偏置是否匹配否则模型会直接报维度错误或者表面上能跑但精度异常低。3.2 训练超参和优化器怎么调Swin 的训练策略和 CNN 有一套明显不同的路子。官方在 ImageNet-1K 上训练 Swin-T 用的是 300 epoch、AdamW 优化器、初始学习率 1e-3、weight decay 0.05、batch size 1024还有 20 epoch 的线性 warmup 和 cosine 学习率衰减。这个配置到我们自己数据集上不能照搬尤其是 batch size。我踩过最深的坑是直接沿用 ResNet 训练时的学习率 0.1Swin 上去以后 loss 直接飞了完全收敛不了。后来换成 AdamW初始学习率按 batch size 缩放比如单卡 batch size 128 时用 2.5e-4多卡 8 卡 batch size 1024 再回到 1e-3才恢复正常。另外 warmup 一定要留够我的经验是最少 5 个 epoch不然早期纯随机状态下相对位置偏置会被带偏后面很难拉回来。数据增强方面Swin 对 RandAugment 和 Mixup 都比较敏感。如果你用的是轻量级训练流程开了太强的增强反而会掉点。我一般会从官方默认的增强强度往下调一点同时把 label smoothing 设为 0.1这个组合在大多数中小数据集上表现比较稳。混合精度训练建议直接用Swin 在 FP16 下稳定性尚可但要注意 loss scalar 的更新策略如果是自定义训练循环最好打开GradScaler的动态调整否则早期收敛会比较慢。3.3 下游任务改造与特征金字塔对齐Swin 在检测分割任务里能无缝衔接主要靠的是它天然输出四层金字塔特征。在 Detectron2 或 MMDetection 里使用 Swin backbone一般流程是把四层特征分别送到 FPN对应 P2 到 P5再接检测头。Swin 的四个 stage 输出分辨率分别是 H/4、H/8、H/16、H/32和 ResNet 系列完全一致所以很多现成框架改动很小。但有一个特别容易踩的细节输入图片的尺寸必须能被 32 整除。Swin 的 patch merging 会对特征图做 2x2 合并如果输入长宽不是 2 的幂次倍数到了后面 stage 特征图尺寸会出现小数导致 reshape 报错或特征对齐错位。我拿自定义数据集做检测时就因为在数据预处理里漏掉了 resize 到 32 的倍数结果小目标漏检非常严重排查了整整一天才发现是这个问题而不是模型本身的问题。如果你的任务不是检测而是语义分割推荐用 UperNet 或者 SegFormer 的解码头配 Swin backbone前者和 Swin 的层级特征配合非常成熟后者简单高效改动量也小。分割任务里一般会使用更大分辨率输入比如 512x512 或 1024x1024这时显存压力主要来自解码头backbone 本身因为有窗口机制反而不算太夸张可以先减少解码头的通道数来平衡显存占用。4. 避坑指南我踩过的那些Swin的坑4.1 小数据集微调的正确姿势如果你手里的数据只有几千张直接拿预训练 Swin 端到端微调很容易过拟合。我的做法是先冻结前两个 stage只训练后两个 stage 和分类头等 loss 稳步下降后再逐步解冻。原因很简单前两个 stage 学到的是低层边缘和纹理特征这些特征在不同数据集上都很通用没必要一开始就打乱。解冻的顺序也有讲究先解冻倒数第二个 stage观察几个 epoch 后 val loss 不再下降再解冻倒数第三个。全量微调时drop path rate 要适当调大一点我一般会从 0.1 涨到 0.3 左右这对防止小数据过拟合很有效。还有一个小技巧Swin 的 LayerNorm 在微调阶段不要冻结否则 BN 和 LN 的统计量会对后期学习率变化非常敏感导致最终精度波动很大。4.2 显存不够时优先改输入还是改窗口很多人一碰到显存不足第一反应是减小窗口大小 M比如从 7 改到 5。这个做法我不推荐因为 M 直接决定了相对位置偏置表的维度改小之后预训练权重里的 B 是 (2x7-1)^2 维模型需要插值才能匹配而且效果不可控。更稳妥的方式是优先降低输入分辨率比如从 224 降到 192或者减少 batch size配合梯度累积来保持有效 batch size 不变。如果实在要改 M建议用双线性插值把相对位置偏置表 resize 到新尺寸同时在加载权重后冻住前两个 stage 做一小段 warmup让模型重新适应新的位置表征。我在一个医学影像分割项目里试过把 M 从 7 改成 12输入分辨率提到 512用插值初始化后模型大概多训了 5 个 epoch 才追回原来的精度操作可行但成本不小一般场景不推荐。4.3 Swin 和 CNN 模型怎么选Swin 在精度上限上确实高但也不是适合所有场景。纯粹的 CNN比如 ResNet、ConvNeXt在推理速度和工程生态上还是碾压级优势。Swin 的 FLOPs 虽然看起来和同体积 CNN 差不多但实际吞吐量往往更低尤其在没有深度优化的推理引擎上差距可能达到 20% 到 30%。如果你做的是移动端实时应用老老实实选 CNN 或者混合结构。如果任务是离线训练追求最高精度尤其是细粒度分类、小目标检测、复杂语义分割Swin 通常是首选。另一个经验法则是看你的数据分布和 ImageNet 差异有多大如果数据分布很接近比如自然图像分类Swin 直接微调就很强如果数据是医学影像、卫星遥感这类分布差异极大的场景Swin 反而不如 ResNet 稳定这时候先跑通 CNN 基线再做模型对比更靠谱。5. 从Swin到Swin v2后续发展与选型思路5.1 Swin Transformer为什么能成为行业标杆Swin 之所以伟大不只是因为它拿了最佳论文而是它给出了一个切实可用的视觉 Transformer 范式。ViT 证明了 Transformer 能用在视觉上Swin 则证明了它能高效、稳定地用在各种密集预测任务上。它的层级结构、窗口注意力、移位窗口这三点互相配合几乎成了后来很多视觉 Transformer 的默认配置。之后一年里ConvNeXt 把 ResNet 逐步改成 Swin 的宏观设计结果性能追平了 SwinFocal Transformer 在窗口注意力上加了粗细粒度交互CSwin 引入了十字形窗口。可以说Swin 之后的很多工作都是在它的框架里做局部打磨很少再有人敢轻易推翻整体范式这本身就说明它的设计足够扎实。5.2 常用变体和预训练模型盘点真正做项目的时候你可能会遇到几个重要变体。第一个是 Swin v2它在 Swin 基础上解决了大规模训练不稳定的问题引入了 LayerScale 和新的相对位置编码支持更大的窗口和更高分辨率预训练模型在 ImageNet-22K 和更多数据上都有明显提升。第二个是 Video Swin把窗口注意力拓展到时空维度用于视频分类和动作检测。第三个是 Swin-Unet专门用来做医学图像分割把 Swin 的层级特征和 U-Net 的跳跃连接结合性能比纯 CNN 的 U-Net 好不少。选预训练模型时我一般优先看 Swin v2 的权重因为它通常比第一代更稳。不过要注意输入分辨率Swin v2 有些模型是 256 输入的有些是 384切换时要先确认相对位置偏置是否匹配。Swin-Transformer-OCR 这一类专用变体如果不是对应场景就没必要用了直接用通用预训练模型再微调反而更省心。5.3 项目实战中的选型建议如果项目是目标检测、实例分割、语义分割这类像素级任务我个人的建议是优先把 Swin-T 或 Swin-S 作为 baseline和 ResNet-50 对比一下。通常 Swin-S 在精度上能明显超越 ResNet-50而训练成本并没有高太多。如果显存紧张还可以考虑在 Swin-T 上做蒸馏把精度迁移到轻量模型里部署压力会小很多。如果是移动端实时场景纯 Swin 就别硬上了换轻量 CNN 或者 MobileViT 这类混合模型更合适。如果你的数据量很大尤其是超过百万级Swin-L 甚至 Swin-v2-L 这种大模型会带来更高的上限但训练成本也很高需要足够的分布式训练资源支撑。总之选型不能只看榜单还要结合推理速度、显存预算、数据规模一起权衡。最后分享一个实操细节我最早用 Swin 做目标检测时在 FPN 前没把最后一个 stage 的步长对齐导致小目标漏检特别严重后来才发现是输入尺寸没有按 32 的倍数裁剪特征对齐错位。这个点官方文档写得不明显但实际影响非常大。Swin 的优势是结构优雅、效果上限高但真正要用好它必须在输入分辨率、预训练权重和训练策略上多花点心思。希望这篇拆解能帮你少走一点弯路。