★ FEATURED ARTICLE 1. 先搞清楚 SK Attention 要解决的是哪类麻烦做视觉模型做久了总会撞上同一个坎卷积核到底选多大。3×3 拿局部纹理准5×5、7×7 抓大结构稳但一旦固定下来模型在整个网络里就只能用一种视野看世界。感受野这件事在分类任务里可能还忍得了放到细粒度识别、遥感小目标、医学影像分割这类场景同一个物体在不同尺度下形态差异极大单一尺度的卷积核很容易顾此失彼。SK AttentionSelective Kernel Attention选择性核注意力就是冲这个痛点来的它让网络在每一层、每个通道上自己决定我现在更想看大一点还是小一点。这个模块最早随 Selective Kernel Networks 一起被提出来核心主张很朴素——不要人为指定感受野让特征自己投票。落地到 PyTorch 里它就是一个二十来行、即插即用的nn.Module输入输出形状完全一致可以直接替换主干网络里的普通卷积也能挂在残差分支上做旁路增强。这篇文章面向的是已经能写基础 PyTorch 代码、想把这个模块真正塞进自己模型里跑起来的读者也照顾刚接触注意力机制、想弄明白为什么权重要在分支维度上做 softmax的新手。我会把原理、参数取舍、完整代码、嵌入主干网络的改法、实测现象和踩坑记录一次性铺开代码可以直接抄走用。2. 三步拆解Split、Fuse、Select 到底在干什么2.1 从一次感受野翻车说起早些年我在做一个偏工业质检的项目图像里同时存在几毫米的划痕和几十厘米的结构性缺陷。骨干网络用的是标准残差结构3×3 卷积堆到底。结果很典型小划痕召回还行大缺陷的边界老是碎的模型像是贴得太近看缺一个退后一步的整体视角。当时的土办法是并行加一路空洞卷积然后把两路特征 concat 起来再接 1×1 降维——能涨点但涨得不稳定而且 concat 之后通道数翻倍后面那层 1×1 得重新学怎么分配权重训练前期震荡明显。SK Attention 提供的是另一条路不是简单地把多尺度特征拼在一起让后续层自己消化而是在模块内部就把每个通道该信哪个尺度这件事算出来直接做加权求和。这样输出通道数不变后续层不用改结构多尺度的信息在模块内部就完成了竞争与筛选。这个设计思路上的差别是理解整个模块的钥匙——拼接是把选择权丢给下游加权是在当下就把选择做完。2.2 Split-Fuse-Select三步走的内在逻辑整个模块可以拆成三个动作我个人习惯按分头看、汇总想、分别选来记Split分头看用 M 个不同卷积核尺寸的分支并行处理同一份输入每个分支输出通道数保持一致。这一步的本质是构造 M 种不同感受野的观察结果它们互不干扰各自完整。Fuse汇总想把 M 个分支的输出逐元素相加得到一个融合特征再做全局平均池化压成每个通道一个标量。这一步的目的是拿到一个全局视角下的通道摘要让后面的选择网络知道当前这层每个通道的整体响应强度。Select分别选用一个轻量全连接网络把通道摘要压到低维再分出 M 个头每个头输出 C 维向量把这 M 个向量在分支维度上做 softmax得到每个通道在 M 个分支上的权重且同一通道的 M 个权重之和为 1。最后按权重加权求和输出最终特征。这套流程里最容易被忽略的是 Fuse 那一步的信息压缩。为什么一定要先做全局平均池化因为选择权重的计算必须是全局的、与空间位置无关的——同一个通道在图像左上角和右下角应该用同一个尺度偏好否则权重图会变成一张噪声图。这一点和 SE 模块的做法一脉相承但 SK 用它来做的是分支选择而不是通道重标定。2.3 和 SE、CBAM 摆在一起看差异很多人第一次看 SK 的代码会觉得这不就是 SE 改了几行。表面像骨子里完全不同。我做了一张对照表这是我当初搞混时自己整理的对比项SECBAMSK Attention注意力维度通道通道 空间分支多尺度核维度归一化方式sigmoid各通道独立sigmoid独立softmax分支间竞争是否有额外卷积分支无无有M 个权重含义这个通道多重要这里多重要这个通道该用哪个尺度参数量增量很小约 2C²/r小中等取决于 M 和核尺寸计算开销几乎可忽略小随 M 线性增长关键在于归一化的位置。SE 用 sigmoid每个通道的权重是独立的可以同时全部接近 1也可以全部接近 0它表达的是重要性。SK 用 softmax 且作用在分支维度上同一通道的 M 个权重被强制归一此消彼长它表达的是偏好——预算只有 1分给谁多一点。这个差别决定了 SK 天然适合做多尺度融合而 SE 适合做通道筛选。你在自己的场景里选哪个先问清楚自己缺的是通道重要性判断还是尺度选择能力。3. 原理深挖每一步的数学细节和参数含义3.1 Split 阶段分支怎么搭才不亏Split 阶段的构造自由度其实挺大。最直接的做法是每个分支一个标准 k×k 卷积k 取 3、5、7 之类。官方实现里还引入了分组卷积的groups参数让分支可以走分组形式降低开销。需要注意的是所有分支的输出通道数必须相同因为后面要逐元素相加。这一点和 Inception 那种分支通道数不同然后 concat的结构是反的很多人第一版代码写错就错在这里——照着 Inception 的习惯给每个分支设了不同的 out_channels结果 Fuse 阶段直接形状不匹配。另一个细节是 padding。要想让 M 个分支的输出空间尺寸一致每个分支的 padding 必须等于k // 2假设 stride1。如果哪个分支 stride 不是 1那整个模块的空间尺寸就会变化输出和输入对不上残差连接也就废了。我在实际使用中基本只让 SK 工作在 stride1 的位置下采样的活儿交给前面的卷积或者池化层干这样最省心。3.2 Fuse 阶段为什么是相加而不是拼接逐元素相加这一步从信息论角度看是有损的但它带来的好处远大于损失。相加之后每个通道只保留一份 C 维的融合表示后续的全连接层参数量是 O(C·d) 级别如果改成拼接融合特征变成 M·C 维全连接层参数直接翻 M 倍而且拼接带来的通道顺序信息哪一段来自哪个分支需要网络额外花力气去学这部分开销纯属浪费。更本质的一点相加操作在数学上对应同一空间位置上不同感受野响应的叠加它保留了这里有东西这个信号同时把是哪个尺度看到的这个信息交给了后面的 Select 阶段重新计算。换句话说Fuse 只负责回答这个通道整体活不活跃不负责回答活跃来自哪个尺度。职责分离得很干净这也是我在自己项目里愿意用它的原因——调试的时候每个阶段的中间张量含义都很明确出问题好定位。3.3 Select 阶段softmax 为什么必须加在分支维度这是整个模块最核心、也最容易写错的一行。选择网络的流程是全局平均池化得到s ∈ R^C一个线性层降维z W·sz ∈ R^d其中d max(C/r, L)M 个线性层分别把z映射回 C 维得到 M 个R^C向量把这 M 个向量堆成(M, C)在 M 这个维度上做 softmax用得到的(M, C)权重对 M 个分支输出加权求和。第 4 步如果写成在 C 维度上 softmax整个模块就从选择尺度退化成了通道归一化效果会明显变差而且不容易从 loss 上看出来因为它照样能训练、照样收敛只是涨点消失了。我见过不止一个实现栽在这包括我自己的第一版。判别方法很简单打印权重张量检查weights.sum(dim0)是不是全 1 的 C 维向量。如果是说明 softmax 维度对了如果weights.sum(dim1)全 1那就是写反了。3.4 M、G、r、L 四个参数怎么定这四个参数的实际含义和取值建议我整理成了下面这张表参数含义常见取值取值影响Mkernels 长度分支数量2 ~ 4越多尺度越丰富显存和耗时线性增长kernels各分支卷积核尺寸[3,5] 或 [3,5,7]奇数、递增覆盖不同感受野Ggroups各分支的分组卷积数1G1 可降参但会削弱分支表达能力rreduction降维比例8 ~ 16越小瓶颈越宽参数量越大L瓶颈最小维度下限32防止浅层 C 很小时 d 被压到 1、2 维关于 M 的取舍我的经验是通道数在 64 以下的时候M2、kernels[3,5] 就够用了再加分支收益递减非常明显通道数 256 以上时 M3 才有明显收益。r 一般从 16 起步如果模型本来就小、C 只有 32那就得靠 L 兜底否则C/r会变成 2瓶颈层信息容量太低选择网络的判断会变得很粗糙。这些不是绝对法则但按这个起点调能少走不少弯路。4. PyTorch 从零实现一个可直接复用的 SK 模块4.1 完整代码下面这版是我自己长期在用的实现做了几处小改动让代码更好读也顺手处理了几个边界情况import torch import torch.nn as nn from collections import OrderedDict class SKAttention(nn.Module): def __init__(self, channel512, kernelsNone, reduction16, group1, L32): super().__init__() kernels kernels or [3, 5] self.kernels kernels self.M len(kernels) # 瓶颈维度L 兜底防止浅层通道数过小时 d 退化 self.d max(L, channel // reduction) # Split多分支卷积输出通道数必须一致 self.convs nn.ModuleList() for k in kernels: self.convs.append(nn.Sequential(OrderedDict([ (conv, nn.Conv2d(channel, channel, kernel_sizek, paddingk // 2, groupsgroup, biasFalse)), (bn, nn.BatchNorm2d(channel)), (relu, nn.ReLU(inplaceTrue)), ]))) # Fuse Select共享降维 M 个升维头 self.fc nn.Linear(channel, self.d, biasFalse) self.fcs nn.ModuleList([nn.Linear(self.d, channel) for _ in range(self.M)]) self.softmax nn.Softmax(dim1) # 注意作用在分支维度 def forward(self, x): bs, c, h, w x.shape # Split feats torch.stack([conv(x) for conv in self.convs], dim1) # (bs, M, c, h, w) # Fuse U feats.sum(dim1) # (bs, c, h, w) S U.mean(dim(2, 3)) # (bs, c) Z self.fc(S) # (bs, d) # Select weights torch.stack([fc(Z) for fc in self.fcs], dim1) # (bs, M, c) weights self.softmax(weights).unsqueeze(-1).unsqueeze(-1) # (bs, M, c, 1, 1) out (feats * weights).sum(dim1) # (bs, c, h, w) return out这版和网上流传的早期实现相比主要差别在于用torch.stack统一了张量布局省掉了一堆view和cat出错概率低很多。特别是权重那一段stack(dim1)之后形状天然就是(bs, M, c)softmax(dim1)作用在 M 上几乎不可能写错维度。早期版本用cat拼成(bs, M*c, 1, 1)再view回来中间那步 reshape 一旦写错顺序权重组会串位而且不会报错只会悄悄掉点。4.2 逐段拆解几个容易被忽略的实现细节第一个细节biasFalse。卷积分支后面紧跟 BNbias 会被 BN 的减均值操作完全吸收留着只是白占参数、白占显存。这个习惯我从写第一个网络就养成了SK 模块里分支多省下来的量其实不小。第二个细节fc不加 bias但fcs加。共享降维层不加 bias 是因为后面紧跟 BN 之后的那条链路已经充分非线性了bias 意义不大而fcs的输出直接进 softmax加 bias 相当于给每个分支一个先验偏置允许网络在训练初期就倾向于某个尺度。这个偏置是可学的不会锁死但确实能加快前几个 epoch 的收敛。第三个细节relu(inplaceTrue)的位置。有的实现把 ReLU 放在 BN 前面或者干脆不放。我的建议是保持在 BN 之后和主干网络里的 conv-bn-relu 顺序一致。如果省掉 ReLU多分支会退化成纯线性组合尺度选择的非线性判别能力会弱一截。第四个细节输入通道和空间尺寸的假设。这个模块假定stride1、输入是 4 维 NCHW。如果你的输入是 5 维视频、3D 卷积或者通道在前得自己改。我在做时序特征的时候试过直接套用维度对不上后来是先把(N, T, C)reshape 成(N*T, C, 1, 1)再送进去的——HW1的时候全局池化等于恒等映射模块退化成一个纯通道选择器也能跑但意义就变了不如直接用 SE。4.3 嵌进残差主干两种改法和各自的代价第一种改法是替换式。以标准 Bottleneck 为例把里面那个 3×3 卷积直接换成 SK 模块class SKBottleneck(nn.Module): expansion 4 def __init__(self, inplanes, planes, stride1, downsampleNone): super().__init__() self.conv1 nn.Conv2d(inplanes, planes, 1, biasFalse) self.bn1 nn.BatchNorm2d(planes) # 关键只在 stride1 的位置替换下采样位置保留原卷积 if stride 1: self.conv2 SKAttention(planes, kernels[3, 5], reduction16) else: self.conv2 nn.Conv2d(planes, planes, 3, stridestride, padding1, biasFalse) self.bn2 nn.BatchNorm2d(planes) self.conv3 nn.Conv2d(planes, planes * self.expansion, 1, biasFalse) self.bn3 nn.BatchNorm2d(planes * self.expansion) self.relu nn.ReLU(inplaceTrue) self.downsample downsample def forward(self, x): identity x out self.relu(self.bn1(self.conv1(x))) out self.bn2(self.conv2(out)) # SK 内部已含 BNReLU这里再套一层 BN out self.bn3(self.conv3(out)) if self.downsample is not None: identity self.downsample(x) return self.relu(out identity)注意这里我把 SK 内部的 BN 和外部 BN 叠了一层看着冗余实际是有意的SK 内部每个分支各自有 BN加权求和之后分布会变外面再套一层 BN 做统一校准训练稳定性明显好一些。这层 BN 是我调试时加上的去掉之后在 batch size 较小时 loss 抖动会变大。第二种改法是旁路式不动原卷积把 SK 作为一个并行分支加进去最后和主分支相加或拼接。这种改法对原模型侵入最小适合在已经训好的模型上做增量实验代价是参数量和显存都会涨而且如果旁路权重初始化不好前期会拖慢收敛。4.4 轻量化变体把开销压下来的几种写法标准实现里 M 个分支都是 full conv通道数等于输入通道数参数量是 M × C² × k²。C256、M3、k[3,5,7] 的时候光卷积分支就有 256×256×(92549) ≈ 5.4M 参数这个量级在移动端是没法接受的。我用过三种压缩思路方案做法参数节省效果损失深度可分离化各分支 conv 改 depthwise pointwise约 8~9 倍轻微通道数大时几乎无损共享权重 不同 dilation一个 3×3 深度卷积dilation 分别取 1、2、3大幅中等感受野覆盖变稀疏分支剪枝只保留 k3 和 k5 两分支约 1/3小多数场景够用我实际在移动端部署场景里用的是深度可分离化这一版先把输入分成 depthwise 卷积每个分支用不同核尺寸做深度卷积再共用一个 pointwise 恢复通道。这样参数量从 O(C²k²) 降到 O(Ck² C²)C256 的时候能省下接近一个数量级的参数速度也快了不少。代价是分支之间的表达能力变弱了——深度卷积每个通道独立运算跨通道的信息交互被推迟到了 pointwise 之后选择网络拿到的融合特征信息量下降最终涨点幅度会打折。5. 跑起来之后的实测现象与调参记录5.1 显存和速度的大致量级我在一张中端显卡上做过一组对照实验输入固定为1×256×56×56batch 取 32只测单模块前向加反向。结论是SK 模块M3kernels[3,5,7]的耗时大约是同等通道数单个 3×3 卷积的 3.5 到 4 倍显存占用约 2.8 倍。这个倍率比理论上的 3 倍三个分支要高一点多出来的部分是torch.stack产生的中间张量和权重张量的广播乘法。如果把stack改成循环累加能省一点显存但代码可读性下降我最后选了保留stack。需要说明的是这些数字严重依赖硬件、框架版本和具体形状你的环境里可能差很多我只提供量级参考。真正要评估的时候别信任何人的数字自己torch.cuda.synchronize()加time.time()跑一百次取均值十分钟的事。实测里还有一个反直觉的现象M 从 2 增加到 3耗时涨了约 45%但精度提升只有零点几个点再从 3 加到 4耗时涨 30%精度几乎不动。这说明分支数的边际收益衰减得很快。我现在的默认配置是 M2只有在通道数超过 256、且任务本身对多尺度敏感比如密集小目标检测时才会开到 3。5.2 训练曲线上的几个典型现象第一个现象前几个 epochsoftmax 权重几乎均匀。这是正常的。初始化时各分支的 BN 参数接近融合特征对各分支的区分度不明显选择网络的输出接近 0softmax 之后就接近 1/M。随着训练推进权重才慢慢分化。我做可视化的时候看到一般到总 epoch 的 15% 到 20% 左右权重分布才开始出现明显的双峰。如果你的权重到训练结束还是均匀分布那基本可以判定模块没起作用回去检查 softmax 维度。第二个现象加了 SK 之后前期 loss 会略高于 baseline。这个不用慌。多分支带来了额外的参数和更复杂的梯度路径前期收敛慢一点是正常的。我观察到的规律是baseline 和 SK 版本的 loss 曲线会在总训练量的 30% 左右交叉之后 SK 版本持续领先。如果你的训练轮数本来就少比如只训 30 个 epoch加 SK 可能反而看不到收益这时候要么加长训练要么用旁路式改法保留预训练权重。第三个现象BN 的 running_mean 在分支间差异明显。有一次我打印各分支 BN 的统计量发现 k3 那个分支的 running_var 明显小于 k7 的分支。这说明大核分支的响应更平小核分支的响应更尖。这不是 bug反而印证了模块确实在做尺度区分——如果所有分支的统计量都一模一样说明它们在学同一个东西那就白搭了。5.3 初始化和学习率的两个小心得SK 模块里的fcs层我习惯用小一点的标准差初始化比如std0.01的正态分布而不是 PyTorch 默认的均匀初始化。原因是这层输出直接进 softmax初始化幅度大的话softmax 会立刻饱和到一个接近 one-hot 的分布等于开局就锁死了分支选择后面很难翻转。用小初始化让权重起步接近均匀给网络留出探索空间。学习率方面加 SK 之后我一般会整体降 10% 到 20%。多分支结构对学习率更敏感尤其在使用较大 batch 和较大基础学习率的配置下不降的话前几个 epoch 很容易出现 loss 尖峰运气不好还会把 BN 的统计量带偏后面怎么训都缓不过来。这算是我交过学费的一条经验。6. 常见问题排查从报错到掉点6.1 形状和维度类报错这类问题最容易定位列个速查表就够了报错信息原因处理The size of tensor a must match tensor b各分支输出通道数不一致统一所有分支的 out_channelsExpected 4D input输入不是 NCHW检查是否误传了 (N, C) 的特征向量mat1 and mat2 shapes cannot be multipliedfc输入维度与 C 不符确认channel参数和实际输入通道一致softmax后权重全为 1/M 且不变化softmax 维度错或初始化过大检查dim调整fcs初始化输出尺寸和输入不一致某分支 padding 或 stride 写错padding 统一为k//2stride 固定 1其中第三条我踩过。当时是在一个通道数会动态变化的模块里直接用 SKchannel参数写死成 256但实际输入是 192报错信息指向的是全连接层看着像是线性层的问题实际是上游通道对不上。后来我在__init__里加了一行断言把这类问题提前到构建阶段暴露调试成本立刻降下来。6.2 能训但掉点几个隐蔽的原因原因一softmax 维度错。前面说过这个错误不报错、不崩只是安静地让模块失效。排查方法在 forward 里临时打印weights.sum(dim1)分支维度和weights.sum(dim2)通道维度前者应该全 1后者一般不等于 1。这一步我建议每个人接手别人代码时都做一遍。原因二模块放在了 stride1 的位置。有些实现为了省事直接在主干的下采样卷积上套 SK结果各分支内部 padding 设计没跟上输出空间尺寸和残差分支对不上表面能跑因为最后做了插值或者形状恰好凑上实际特征已经错位了。我的原则很硬SK 只放 stride1。原因三BN 的 momentum 太大。多分支结构里每个分支的 BN 各自统计自己的分布如果 momentum 设在 0.1 以上batch 又小统计量抖动会很大导致各分支的输出尺度反复变化选择网络的权重也跟着抖最后收敛到一个很差的局部。我一般把 SK 内部 BN 的 momentum 调到 0.05 甚至 0.03配合稍大的 batch 使用。原因四把 SK 加在了已经很强的预训练模型上但没做 warmup。新加的模块是随机初始化的一上来就用正常学习率训会把原来学好的特征冲掉。这个坑我踩过两次后来固定流程是先冻结主干只训新增模块 2 个 epoch再解冻整体用较小的学习率微调。6.3 部署侧的坑导出静态图的时候要注意torch.stack加sum的组合在某些推理框架里会被拆成很多小算子性能比预期的差。我在一次端侧部署里测出来SK 模块在推理引擎里的实际耗时比在 PyTorch 里测的高了近一倍。解决办法是把stack改成显式的 M 次循环累加虽然代码长一点但算子图更扁平推理引擎能更好地融合。另外softmax在某些量化流程里对精度很敏感。动态量化的时候权重张量的数值范围跨度很大有的接近 0有的接近 1按张量整体量化会丢精度。我最后是把 softmax 前后的部分拆开处理softmax 保持浮点只量化前后的线性层和卷积层精度损失控制在可接受范围内。这块没有通用方案得按你的部署栈自己试。7. 什么场景值得上 SK什么场景劝你别碰7.1 值得用的几类任务多尺度目标共存的检测与分割。这是 SK 最擅长的场景。遥感影像、医学影像、工业质检里同一张图里目标尺寸差异经常超过一个数量级固定感受野的模型顾此失彼。SK 在通道维度上给每个尺度分配权重等于让网络自己学这一层该看多大比人工设计特征金字塔省事也比纯粹的空洞卷积组合更灵活。通道数较大的深层网络。SK 的收益和通道数正相关。C 在 256 以上时通道维度的选择空间足够大权重分布能学出比较清晰的模式C 只有 32、64 的时候选择网络的判别能力有限收益经常被随机波动淹没。所以浅层或者窄网络里加 SK期望值要放低。已经有较好 baseline、想再抠几个点的情况。SK 的定位一直是锦上添花型模块它不会把一个烂模型救活但能在一个已经调好的模型上稳定贡献一点提升。如果你的 baseline 还没调好先别急着上注意力模块把数据、增强、学习率策略这些基本功做扎实收益可能比加模块大得多。7.2 建议绕开的场景极端轻量化的移动端模型。参数量和延迟预算卡得很死的时候SK 的多分支结构性价比不高。真要加用前面说的深度可分离变体而且只在高分辨率特征图上加低分辨率层省掉。训练数据量很小的任务。多分支结构意味着更多参数和更复杂的梯度路径数据少的时候过拟合风险明显上升。我在一个只有几千张图的小数据集上试过加了 SK 之后训练集精度涨得很快验证集反而掉了。后来是靠强增强加重度正则才压住但收益也就勉强打平并不划算。推理时延敏感的在线服务。前面测过M3 的时候单模块耗时是普通卷积的三倍多即使轻量化之后也比 SE 这类纯通道注意力重。线上的时候先算清楚你的延迟预算够不够再决定。7.3 几个可以继续往下做的小方向一个是把分支选择做成分层的不是一次在所有分支上做 softmax而是两两分组先选一次、再在胜者之间选一次类似锦标赛结构。这样 M4 时可以用 3 次二分类代替一次四分类参数量下降理论上判别能力还能保持。我做过一版原型训练不太稳但方向上我觉得是通的。另一个是把 SK 和空间注意力结合起来在分支选择之外再叠一个空间权重形成尺度 × 空间的二维选择。思路是在 Fuse 之后不做全局池化而是保留空间维度做卷积降维让不同区域用不同尺度。这个代价是参数量会上去一些而且失去了全局性带来的稳定性需要更多的正则。最后分享一个我自己在用的调试小工具写一个 hook 挂在 SK 模块上每个 epoch 记录一次权重张量按通道求均值的分布画成直方图。如果分布是单峰的、贴近 1/M说明模块没学到东西如果呈现清晰的双峰说明尺度选择确实在起作用如果分布随时间剧烈震荡说明学习率或者 BN 参数有问题。就靠这个小工具我排查过三次以为在涨点其实模块是死的情况比看 loss 曲线灵敏得多。
阅读完成 · 觉得有帮助？
