之前有一阵子我一直在琢磨一件事ViT出来以后图像分类确实被打通了可一旦把它接到检测、分割这类密集预测任务上总觉得哪里别扭。同样是主干网络ResNet的脖子和头用得很顺ViT反而有种“有力气使不出来”的感觉。后来我把Swin Transformer这篇2021年初的工作从头到尾过了一遍才意识到问题不在Transformer本身而在于ViT那种“一路保持单分辨率、不做金字塔”的设计跟CV下游任务的需求压根不对齐。Swin Transformer用移位窗口和分层结构把Transformer做成了CNN一样具备多尺度特征图的主干正好补齐了这块拼图。这篇博文我会从动机、架构、核心机制到代码实现细节把Swin Transformer拆开讲清楚尤其会重点说patch merging是怎么一步步把分辨率降下来、通道数提上去最终输出一组多尺度特征图的。适合正在读论文、准备复现实验、或者想把Transformer主干接进检测分割项目的朋友参考。1. ViT的天花板为什么整图注意力在CV任务上“使不出劲”1.1 ViT把CV问题简化成了“大块头分类问题”ViT的思路很直接把一张224×224的图像切成16×16的patch每个patch线性投影成一个token然后像处理句子一样让所有token在全序列长度上做全局自注意力。这个做法在ImageNet分类上表现很好但它隐含了一个前提你只关心“这张图是什么”不关心“东西在哪、边界在哪、每个像素属于谁”。换句话说ViT设计的时候就没把空间分辨率当回事它把空间信息压缩进了patch embedding然后在token层面做全局关系建模。问题恰恰出在这。检测、分割这类任务需要像素级的定位能力需要在多个尺度上观察目标。小目标可能只有十几个像素大目标能占满半个画面。ViT从头到尾都是14×14或者16×16的token网格相当于只有一层尺度没有“由细到粗再回到细”的过程下游任务拿到这张特征图既缺细节也缺语义层次。1.2 全局注意力的计算复杂度是一堵墙再往深一层说即便ViT想做多尺度全局自注意力本身就扛不住高分辨率输入。计算一下标准多头自注意力的复杂度就知道了。输入token数L特征维C复杂度大概是O(L²C LC²)如果输入是224×224patch是4×4那L就是56×563136。平方一下约千万级别的计算量已经有点吃力了。如果输入换成1024×1024L变成256×25665536L²直接到43亿这谁顶得住所以Swin的思路很务实既然全局注意力算不动那就别硬算。把注意力限制在局部窗口里窗口之间再通过移位机制交换信息。这样既保留了Transformer建模长程依赖的能力又把复杂度从O(L²)降到了跟token数线性相关的水平。1.3 CNN的“分层”其实是CV任务最朴素的先验还有一个点很关键CNN的多尺度结构不是偶然的它是被验证了几十年、极其有效的视觉先验。浅层抓边缘纹理中层抓部件深层抓整个物体这本身就是人类视觉系统的工作方式。FPN、UNet、各类检测分割框架全都建立在“多尺度特征图”这个基础上。Swin Transformer的高明之处就是没有跟这个先验对着干。它没有试图用Transformer替代整个CNN而是把Transformer塞进了CNN的分层骨架里。每个stage内部是局部窗口注意力stage之间用patch merging降采样输出的特征图天然就是FPN需要的多尺度金字塔。这个设计取舍决定了后来Swin能在检测分割任务上全面压制同体量的CNN主干也成为它跟ViT系列最大的分水岭。2. 数据流全解patch merging如何一步步打造多尺度金字塔2.1 第一步4×4 patch切分先把分辨率降四倍Swin的输入处理跟ViT就不一样。ViT一般用16×16的patchSwin用4×4。这个选择不是拍脑袋而是为了保留更多空间细节给后续逐级降采样留出余地。具体流程是输入图像H×W×3切分成(H/4)×(W/4)个patch每个patch是4×4×348维经过一个线性embedding层投影到C维。论文里Swin-T的C96所以经过Patch Partition Linear Embedding之后得到的是(H/4)×(W/4)×96的特征图。这里有个容易忽略的细节Patch Partition在代码里通常直接用卷积实现。一个kernel_size4, stride4的卷积输出通道设为C一步就能完成“切patch 线性投影”两件事。这也是PyTorch官方实现的标准做法。很多人初读论文时以为要先切patch再做embedding实际上一个卷积全搞定了。2.2 Stage 1保持分辨率做局部Transformer建模拿到(H/4)×(W/4)×96的特征图进入Stage 1。Stage 1由若干个Swin Transformer Block组成Swin-T在这里放了2个Swin-B放了2个。这个阶段的关键是分辨率不再下降通道数也保持不变纯粹做特征建模。每个block内部是窗口注意力 MLP窗口大小M默认是7。以56×56的特征图为例窗口数就是(56/7)²64个每个窗口里7×749个token。到此为止Swin还只是在“第一层尺度”上做文章。真正让它具备金字塔能力的是后面的patch merging。2.3 Patch Merging间隔采样 通道拼接 线性映射Stage 2、3、4的入口都是Patch Merging。它是Swin的分层关键作用是把空间分辨率降一半、通道数翻一倍完全对标CNN里的stride-2下采样。具体做法分三步对输入特征图做间隔采样。取位置(x, y)如果x和y都是偶数就单独抽出来一组x偶y奇抽一组x奇y偶抽一组x和y都奇数再抽一组。这四组分别对应“左上、右上、左下、右下”的像素每组形状都是(H/2)×(W/2)×C。把四组特征在通道维度上拼接得到(H/2)×(W/2)×4C。通过一个线性层映射回2C维再经过LayerNorm。经过这三步输入是56×56×96输出就变成28×28×192。分辨率减半、通道翻倍信息没有丢只是重新组织了一下。这里我提醒一句这个“间隔采样”在代码里不是用循环做的而是先用一个reshape把特征图拆成四个方向上的offset再直接concatenate。如果你自己复现别写成四层for循环效率差太远。2.4 Stage 2/3/4分辨率逐级减半通道逐级翻倍每个Stage的结构都是“Patch Merging 若干Swin Block”Stage 2先合并到14×14×384以Swin-T为例再经过2个blockStage 3先合并到7×7×768再经过18个block——这是Swin-T最深的stage计算主力都在这Stage 4这里有个细节Stage 4之前还会再做一次Patch Merging把特征图压到7×7? 等一下我重新算一下。以224×224输入、C96的Swin-T为例Patch Partition后56×56×96Stage 1不改分辨率56×56×96第一个Patch Merging28×28×192Stage 228×28×192第二个Patch Merging14×14×384Stage 314×14×384第三个Patch Merging7×7×768Stage 47×7×768所以Stage 4的输出是7×7×768分辨率从224一路降到7刚好是32倍下采样跟ResNet最后一级输出的下采样率完全一致。而中间每一级都会留一份特征图出来这就是现成的多尺度特征金字塔阶段分辨率224输入通道数Swin-T下采样倍数Stage 156×56964×Stage 228×281928×Stage 314×1438416×Stage 47×776832×这组特征往FPN里一塞检测分割直接拿来用。这也是Swin能无缝对接检测框架的根本原因——它的输出格式本来就是多尺度的不需要像ViT那样先把14×14的token网格拉直再强行reshape回一层特征图。3. 移位窗口注意力把全局注意力关进窗口再让窗口自己串门3.1 窗口注意力的复杂度优势是怎么来的标准多头自注意力里每个token都要跟全图所有token做attention。Swin把特征图切成M×M的窗口窗口内部做自注意力。以Swin-T的Stage 1为例56×56的特征图切成64个7×7窗口。每个窗口内部49个token做自注意力计算量是64×49²153K量级。如果做全局注意力是3136²≈9830K差了64倍。公式化一点标准注意力复杂度是O(L²C)窗口注意力是O(WM²C)O((L/M²)·M²·C)O(LC)。注意L在指数上变成了1次方这是从“二次复杂度”到“线性复杂度”的本质转变。分辨率越大这种省下来的算力越可观。3.2 固定窗口的缺陷信息被框死在局部但代价也很明显窗口之间没有交互。左上角的窗口永远看不到右上角窗口里的内容整个特征图被切成了互不通信的小块。直白点说纯窗口注意力只能建模“局部关系”跟卷积的感受野拉不开差距这不就失去了用Transformer的意义吗Swin的解法是奇数层用规则窗口W-MSA偶数层把窗口整体往右下方向偏移半个窗口大小SW-MSA。两次交替之后原本处于不同窗口区域的token在新的划分方式下会被分到同一个窗口彼此之间就能直接做注意力了。打个比方第1层里a和b隔了一个窗口谁也看不见谁。第2层窗口往右下方挪了半格a和b可能就被分进同一个新窗口。它们的特征第一次产生了直接交互。逐层堆叠下去信息就能跨窗口一步步传递最终覆盖全局。3.3 循环移位和掩码高效实现的关键移位窗口听上去简单实现起来有个麻烦如果窗口整体偏移了半个M特征图边缘会多出一圈窗口数量会从规则的4个变成9个而且窗口尺寸不一致没法直接批量计算。Swin用的是循环移位cyclic shift来处理这个问题。把特征图沿左上方向把多出来的部分裁剪下来拼到另一侧去让特征图重新变成一个可以被规则窗口切分的形状。这样窗口数量保持不变尺寸也一致可以并行计算。但这么一挪有些窗口内部混入了原本不相邻的区域。比如一个窗口里左上角那块是原图的右下角右下角那块是原图的左上角它们物理距离很远本来不应该做注意力。Swin为每个窗口生成一个掩码attention mask把不属于同一物理区域的token对在计算attention score时直接加一个很大的负数softmax之后贡献归零。很多人读论文时觉得“移位窗口”好理解一到掩码就头疼。其实核心只有一句话窗口位置变了但每个token知道自己原本属于哪个物理区域注意力只允许在同一个物理区域内发生。3.4 为什么说“交替使用”是灵魂我刚才讲的W-MSA和SW-MSA交替不是可有可无的锦上添花它是整个Swin架构的骨架逻辑。如果只有W-MSA等于用一堆互不通信的局部Transformer块做特征提取效果大概率不如同规模的CNN因为卷积至少还有3×3邻域的重叠感受野。如果只有SW-MSA那每次都要循环移位和掩码计算效率低而且失去了局部归纳偏置的稳定性。交替使用的好处是第奇数次是规则窗口计算高效第偶数次是移位窗口促成跨窗口信息流动。两个block各司其职合在一起才是一个完整的“局部建模 全局通信”单元。这也是Swin Block数量在Swin-T里是2的倍数2、2、18、2的原因——为了保持这种成对的结构。4. 相对位置偏置Swin精度上涨的隐形功臣4.1 绝对位置编码在窗口注意力里的水土不服ViT用的是一个可学习的绝对位置编码向量加到patch embedding上。Swin不能直接照搬原因在于窗口注意力窗口在图中不断移动同一个窗口内token的绝对位置是变化的共享同一套注意力参数时绝对位置带来的偏差会干扰原本应该纯粹由内容决定的attention_score。Swin选择的是相对位置偏置relative position bias。每个注意力头在计算attention时会在query和key的点积结果上加一个只取决于两者相对位置的偏置项。Attention(Q, K, V) SoftMax(QKᵀ / √d B)V其中B就是相对位置偏置。4.2 偏置表的形状为什么是(2M−1)×(2M−1)窗口大小M7窗口内任意两个token的坐标差范围是−6到6。两个方向的差值各自有13种可能−6到6因此所有可能的相对位置组合是13×13169。这就是论文里那个偏置表B的维度(2M−1)×(2M−1)的来历。Swin-T在Stage 1用M7所以每层每个注意力头对应一个13×13的可学习矩阵。实现时有个细节偏置表的索引不是两两坐标差的绝对值而是带方向的坐标差。一个token在另一个的右上角和它在另一个的左下角相对位置偏置是不一样的。这让模型能感知方向性对物体朝向、几何结构的建模有明显帮助。很多复现代码里会先算出一个“相对位置索引表”形状是M²×M²每个元素是二维坐标差。然后把这个索引表转成一维的扁平索引去查B表。这一步在代码实现里特别容易出错——坐标差做算术运算后要确保映射到非负区间而且索引顺序要跟attention score矩阵的行列对应。给你一个简单的验证方法写代码的时候把B表设成可学习的并打印出来看看有没有梯度更新。如果B表在训练中一直是随机初始值八成是索引查表写错了。4.3 共享与复用的哲学还有一点值得提Swin的位置偏置表不是每层各自独立的而是每层内部所有注意力头共享一个刚查出来的相对位置索引表但偏置数值B不同。这意味着同一层里不同窗口可以复用同一套相对位置关系——不管窗口在图中的哪个位置“左移一格、上移一格”的相对关系都一样。这个设计非常优雅。它让模型学到了平移不变的相对位置先验又保留了不同层、不同头各自微调位置语义的自由度。论文里专门做了消融去掉相对位置偏置精度掉一大截换成绝对位置编码精度也不如带偏置的版本。这个不起眼的B表实际是Swin精度能追平甚至反超CNN的关键因素之一。5. 三个关键的工程细节看起来小错起来大这一节不讲概念专门讲我复现和实践过程中踩过的、以及观察别人踩过的几个细节坑。5.1 Patch Merging里的reshape与permute顺序Patch Merging的间隔采样在代码里是用reshape permute实现的不是循环。第一次写的时候我按照直觉先reshape成(B, H/2, 2, W/2, 2, C)然后直接取对角线元素结果输出的空间顺序全乱了。正确做法是先permute成(B, H/2, W/2, 2, 2, C)让两个2分别代表行方向和列方向的采样偏移再reshape成(B, H/2, W/2, 4C)。这一步的教训是先想清楚维度语义再写permute。resize的索引顺序、采样的空间映射关系稍微差一点后面stage的特征就对不齐了。5.2 掩码计算在窗口不整除时怎么办标准输入是224×224patch partition后56×56窗口7×7刚好整除所以很多人忽略了窗口切分可能遇到的边缘残缺问题。但是只要输入分辨率变一下比如测试时用480×480patch后变成120×120120除以717.14切出来的最后一个窗口不够7×7。PyTorch的窗口切分一般会用padding补齐然后把padding区域用掩码遮挡掉。我建议不要依赖“刚好整除”这个巧合把代码里的pad逻辑写完整。至少包括计算H方向需要pad的数量按M的整数倍向上取整。W方向同理。给attention mask留出pad位置设置为−100保证padding区域不参与注意力。很多公开实现的测试推理部分会专门处理这个问题训练时因为固定224就省略了。到了部署阶段一旦输入尺寸变化这个坑必炸。5.3 相对位置索引表的注册方式PyTorch里偏置表通常注册成buffer因为它不是参与梯度更新不它是可学习的参数应该用Parameter注册。但相对位置索引表不含梯度应该用register_buffer。我见过一种错误写法把整个偏置表都用buffer注册训练时偏置完全不更新模型精度怎么都上不去。检查方法很简单print模型参数的requires_grad列表看看位置偏置对应的项是不是False。正确结构是relative_position_indexregister_bufferM²×M²的长整型索引表不更新。relative_position_bias_tablenn.Parameter维度(2M−1)×(2M−1)×num_heads参与更新。6. Swin各尺寸配置、计算量与适用场景的参考6.1 配置表的背后逻辑Swin论文提供了一组不同规模的主干配置参数规模从28M到197M不等模型通道C每层block数窗口大小参数量Swin-T962, 2, 18, 2728MSwin-S962, 18, 18, 2750MSwin-B1282, 2, 18, 2788MSwin-L1922, 48, 24, 27197M需要说明的是block数不对称Stage 3放了最多block因为降到14×14之后计算量已经可控可以放心堆深度。Stage 4只有2个block因为7×7本身token极少不需要继续加深。6.2 计算量的分配也很有意思以Swin-T为例总FLOPs约4.5G其中Stage 3占了大部分。这是因为Stage 3分辨率14×14窗口8×864个每个窗口49个token虽然token数少但block数量多。如果只做分类一般用Stage 4输出的7×7特征图全局池化后接分类头。如果做检测分割就把四个stage的输出全部接进neck。这算是Swin主干的通用使用范式。7. 实际项目里Survive下来的几条经验7.1 训练超参的几个“默认答案”我试过几组配置比较稳定的组合是优化器用AdamW初始学习率1e-3左右weight decay 0.05。学习率调度用cosine退火配合20个epoch左右的warmup。数据增强跟ViT那套ImageNet配置接近RandAugment、Mixup、CutMix、Random Erasing都可以上。分类任务训练300个epochImageNet top-1一般在83%上下Swin-T跟论文报告比较接近。有一点跟CNN明显不同Swin对学习率很敏感。learning rate稍微大一点比如直接上3e-3前面warmup阶段loss就容易飞。建议从小学习率1e-3起步用warmup过渡到1e-3严格的顺序是warmup阶段从0或极小值线性增加到目标学习率再用cosine退火降下来。CNN常用的“常数学习率step decay”在这里不太灵至少我试下来效果差不少。7.2 DropPath我自己踩过的正则化陷阱Swin Block里有个常被忽略的正则化手段——DropPath即随机丢弃整个残差分支。Swin-T用的是0.1的默认值。如果做小数据集比如只有几千张图0.1可能不够建议适度上调到0.2~0.3能缓解过拟合。但如果数据量很大DropPath开大了反而会导致训练不稳定。这里有个容易踩的坑DropPath应该加在block内部残差连接的分支上即x x drop_path(attention(x)) x x drop_path(mlp(x))而不是加在整个block的输出上。两者效果差很多我见过有人把DropPath直接加在block输出端结果整个网络的恒等映射被频繁打断特征退化严重。7.3 下游任务怎么接才顺做检测时把Stage 1到Stage 4的输出按顺序送进FPN注意Stage 1输出的通道数C通常太小96如果直接进FPN可能细节有余语义不足很多框架里会从Stage 2开始接。做分割时UNet式的跳连结构跟Swin的多尺度特征天然契合。我用Swin-T做过一个语义分割demo把Stage 1的高分辨率特征作为skip connectionStage 4的低分辨率语义特征作为bottleneck上采样路径一层层融合回来效果比同参数量的CNN主干好一截。7.4 推理阶段的性能特征Swin窗口注意力虽然计算量低但实现上需要对特征图做窗口重排window partition/unpartition这类操作在GPU上会引入额外的访存开销尤其在小batch、小分辨率下特别明显。实测下来Swin-T在单卡上的推理吞吐比同FLOPs的ResNet略低但在大分辨率输入场景比如检测任务里的1216×1216输入Swin的优势会体现出来——因为窗口注意力的复杂度跟分辨率线性相关而ResNet在深层靠的是stride卷积特征图小了以后计算量下降但没Swin省得这么明显。如果项目对延迟特别敏感建议优先考虑Swin-T并把窗口大小从7调成4试试。窗口变小注意力矩阵从49×49变成16×16计算量大幅下降精度损失看任务有的任务几乎不掉点。8. 写在复现笔记最后的话Swin Transformer这个架构从问世到现在已经成了视觉Transformer分支里一个绕不开的参考点。它是“把Transformer做成分层主干”这件事的标杆patch merging和移位窗口这两个设计已经被后续很多工作吸收和改造。如果让我给正准备读这篇论文的朋友一个建议第一遍读的时候抓住三条主线就好——多尺度是怎么来的、窗口注意力怎么算的、相对位置偏置是干嘛的。第二遍再去看代码细节尤其是掩码和索引表这两块。第三遍如果做复现从Swin-T开始先跑通分类再接到检测分割里看多尺度特征图的实际效果。读论文最怕停在“知道概念”这个层面。Swin的价值不只是提出了一堆名词而是把“多尺度、窗口、偏置”这些零件组装成了一个能直接用的系统。真正想搞懂它跑一个demo比读十遍摘要都管用。
