1. 从标题拆解出的核心命题
第一次看到“多方向融合特征金字塔”这个说法,我脑子里冒出来的第一个念头是:特征金字塔这东西不是已经被玩烂了吗,怎么还能在顶会上单独拎出来讲。后来把论文翻了几遍,又自己动手复现了一版,才明白它真正想解决的是一个很具体、很工程化的痛点——标准特征金字塔在融合不同尺度特征时,信息流动的方向是单一的,要么自顶向下,要么自底向上,横向连接只是简单相加,导致跨尺度语义对齐做得不够充分。
这个项目标题里有两个关键词值得先掰开说清楚。一个是“借鉴Transformer”,另一个是“多方向融合”。前者说的是它把Transformer里那套多头注意力的思路搬到了特征金字塔的构建上,后者说的是它不再满足于单向的路径聚合,而是让特征在多个方向上来回流动、反复对齐。这两件事合在一起,构成了FPT(Feature Pyramid Transformer)的核心骨架。
我先把结论摆前面:FPT想干的事情,是让特征金字塔里的每一个层级的特征,都能“看到”其他所有层级的特征,并且用注意力机制决定该吸收多少、从哪个方向吸收。传统FPN里高层语义往下传、低层细节往上走,中间靠一个1x1卷积降维后直接相加,这种做法在目标检测里够用,但遇到尺度差异大、小目标密集的场景就容易露怯。FPT的思路是把这个“相加”换成“注意力聚合”,而且聚合的方向不止一个。
适合谁来读这篇内容?如果你正在做目标检测、实例分割这类需要多尺度特征融合的任务,或者你手头有个自研的检测头想加点跨尺度建模的能力,再或者你只是好奇Transformer那套注意力机制除了在NLP和ViT里刷榜之外还能怎么用,那这篇拆解应该能给你一些可以直接抄的思路。下面我会从设计动机、核心结构、实操复现、踩坑排查几个层面把它讲透,尽量做到看完就能动手改自己模型里的neck部分。
2. 整体设计思路与方案选型考量
2.1 为什么标准FPN的融合方式不够用
要理解FPT为什么这么设计,得先回到FPN本身的局限上。标准FPN的结构大家都很熟:骨干网络吐出C3、C4、C5几个stage的特征图,然后经过横向连接降维到统一通道数,再自顶向下逐层相加,得到P3、P4、P5。这个过程中,高层特征(语义强、分辨率低)通过上采样和低层特征融合,低层特征获得了语义信息,这是它有效的原因。
但问题在于,这个融合是单向且一次性的。P5的信息传到P3要经过两次上采样和两次相加,每次相加都是简单的逐元素加法,没有任何权重分配机制。换句话说,P3在吸收P5传来的语义时,并不知道P5的哪些空间位置、哪些通道对当前任务更重要,它只是被动地全盘接收。反过来,低层的位置细节往高层传的路径在原始FPN里压根就不存在(PANet后来补了这条自底向上的路径,但依然是简单相加)。
我在实际项目里遇到过很典型的情况:一张图里既有占几十个像素的大目标,又有十几个像素的小目标,标准FPN在P3上做小目标检测时,语义信息是从P5一路加下来的,中间经过两次无差别相加,到了P3已经混入了不少来自大目标区域的噪声响应。这就是单向无权重融合的代价。
2.2 把注意力机制引入金字塔的三个方向
FPT的核心洞察是:既然Transformer里的自注意力能让每个token动态地聚合全局信息,那特征金字塔里的每个空间位置,是不是也可以动态地从其他层级聚合信息?基于这个想法,它设计了三种信息流动方向,这也是“多方向融合”这个说法的由来。
第一种是自顶向下的非局部融合。高层特征作为query,低层特征作为key和value,让高层的每个位置去低层里找与之语义相关的区域。这跟标准FPN的自顶向下不同之处在于,它不是把整张高层特征图无差别地上采样后加过去,而是让高层特征主动“查询”低层特征,按相关性加权聚合。直觉上,高层的一个语义区域(比如一只狗的头部)会去低层里找对应的细节区域(狗的毛发纹理),而不是把整张低层图都平均进来。
第二种是自底向上的非局部融合。反过来,低层特征作为query,高层特征作为key和value。低层的每个位置去高层里找语义上相关的区域,把高层语义按相关性拉下来。这条路径补上了原始FPN缺失的自底向上语义传递,而且同样是有权重、有选择的。
第三种是同层内的非局部融合。这个方向不跨层级,而是在同一层特征图内部做自注意力,让每个位置聚合同层其他位置的信息,增强层内的全局上下文。这一路的作用类似于在每一层内部先做一次全局建模,把空间上的长程依赖先理一遍,再去跨层融合。
注意:这三个方向的融合不是串行堆叠,而是并行处理后做聚合。这一点在复现时很容易搞错,如果按串行方式实现,计算量和显存占用会翻好几倍,而且效果反而下降。
2.3 为什么用Transformer而不是别的聚合方式
有人可能会问,跨尺度融合用注意力,为什么非得是Transformer那套?用SE模块、CBAM这类通道注意力不行吗?我的理解是,通道注意力解决的是“哪些通道重要”,空间注意力解决的是“哪些位置重要”,但它们都是在单一尺度内部做文章。FPT要解决的是跨尺度、跨空间位置的相关性建模,query来自一个尺度、key/value来自另一个尺度,这种异构的注意力计算天然适合用Transformer的QKV范式来表达。
而且Transformer的多头机制在这里有个额外好处:不同的头可以学到不同的跨尺度对应关系。比如一个头专门关注大目标整体到局部细节的映射,另一个头关注小目标区域到高层语义的映射,多个头并行捕捉多种融合模式。这一点在可视化注意力图的时候能明显看出来,不同头的关注区域确实有分工。
从工程角度看,选Transformer还有一个现实考量:它的实现可以复用现成的多头注意力算子,不需要从零写CUDA,在主流框架里都有高效实现,落地成本可控。这也是我当初愿意花时间复现它的原因之一。
3. 核心结构拆解与关键细节
3.1 三个融合分支的具体计算流程
把整体思路理清之后,接下来得把每个分支的计算细节抠清楚,不然复现的时候全是坑。我按自己的理解把三个分支的计算流程重新梳理一遍。
自顶向下分支的处理逻辑是这样的:取高层特征图,经过1x1卷积调整通道后展平成序列,作为query;取低层特征图,同样调整通道后展平,作为key和value。然后做标准的多头交叉注意力,得到聚合后的高层特征序列,再reshape回空间维度。这里有个关键细节:高层特征的分辨率低、序列短,低层特征分辨率高、序列长,所以这个注意力是“短query查长key”,计算量相对可控。聚合完之后,把结果上采样到低层分辨率,与低层特征做融合。
自底向上分支正好反过来:低层特征做query,高层特征做key和value,是“长query查短key”,计算量会大一些。聚合结果下采样到高层分辨率后融合。这里下采样的方式论文里用的是步长卷积或者池化,我实测下来用步长2的3x3卷积比直接池化效果略好,因为卷积能保留一些可学习的信息压缩能力。
同层分支最简单,就是标准的多头自注意力,query、key、value都来自同一层特征。但要注意,如果对每一层都做全分辨率自注意力,P3那层的序列长度会非常夸张(比如输入800x800的图,P3是100x100,序列长度就是10000),显存直接爆炸。所以实际实现里通常会对同层注意力做下采样或者用窗口化的方式限制计算范围,这一点论文里没有特别强调,但复现时绕不开。
3.2 通道维度与序列长度的处理技巧
这里我要重点讲一个复现时最容易翻车的地方:通道维度和序列长度的映射关系。Transformer的注意力是在序列维度上做的,所以必须把特征图从(B, C, H, W)展平成(B, HW, C),注意力计算完再展回去。这个过程中,C就是注意力的embedding维度,HW是序列长度。
问题在于,不同层级的C可能不一样(比如C3是256,C5是512),H*W更是差好几倍。FPT的做法是先用1x1卷积把所有层级的通道统一到一个固定值(论文里用的是256),这样跨层注意力的维度才对得上。序列长度不一致没关系,因为query和key/value的序列长度本来就可以不同,注意力矩阵是(短序列长度, 长序列长度)的形状,这是Transformer天然支持的。
但显存是个硬约束。我实测过,如果对P3做全序列的自注意力,在batch size为2、输入1024x1024的情况下,单卡24G显存直接OOM。解决办法有两个:一是对同层注意力做空间下采样,比如先池化到一半分辨率再做注意力,算完再上采样回去;二是限制注意力的感受野,用局部窗口代替全局。前者实现简单但会损失一些细节,后者更精细但代码复杂。我一般先用前者快速验证效果,确认有收益再考虑上窗口化。
3.3 融合权重的聚合方式选择
三个分支各自算出聚合特征之后,怎么把它们合到一起?论文里给了几种选项,我逐个试过,这里说下实际感受。
最简单的做法是直接相加,三个分支的输出逐元素加起来。优点是实现简单、参数量零增加,缺点是三个分支的贡献被固定成等权,没法自适应。第二种是拼接后过1x1卷积降维,这样网络可以自己学各分支的权重,参数量增加很少,效果比直接相加稳定。第三种是学一组标量权重做加权求和,类似可学习的残差权重,参数量最少但表达能力也最弱。
我自己的项目里最终选的是拼接加1x1卷积的方案。原因很实际:直接相加在训练初期loss下降还行,但收敛后的精度比拼接方案低了大概0.5个点(在自建数据集上的mAP),而拼接方案只多了不到十万个参数,性价比明显更高。标量加权那版我没继续调,因为拼接方案已经够用了,没必要为了省那点参数牺牲精度。
提示:如果你显存紧张,可以只保留自顶向下和自底向上两个分支,去掉同层分支,精度损失大约0.3个点,但显存能省将近四成。这是我做消融时测出来的,供参考。
3.4 位置编码的必要性讨论
Transformer做注意力时,如果不加位置编码,模型对序列里元素的顺序是无感知的。在NLP里这很致命,但在特征图里情况有点不一样——特征图本身是二维空间结构,展平成序列后空间邻接关系丢失了,所以理论上也需要位置编码来补回空间信息。
FPT论文里对位置编码的处理比较轻量,我印象里是在展平后加了一个可学习的位置嵌入。我复现时试过加和不加两种情况,不加位置编码时,同层注意力的效果下降比较明显(因为同层注意力完全依赖空间关系),但跨层注意力的影响相对小一些,因为跨层注意力更多依赖语义相关性而非空间邻近性。
实际落地时我的建议是:位置编码加上,但不用搞太复杂,可学习的一维位置嵌入就够用,成本很低。如果追求更精细的空间建模,可以用二维正弦位置编码,但收益递减,我测下来跟可学习版本的差距在0.1个点以内,不值得为此增加实现复杂度。
4. 实操复现与核心环节实现
4.1 环境准备与骨干网络对接
动手之前先把环境理清楚。我用的是PyTorch 1.10加CUDA 11.3的组合,这个版本组合在我这边跑下来最稳,多头注意力的算子效率也够。骨干网络选的是ResNet50,从torchvision里直接加载预训练权重,取layer2、layer3、layer4的输出作为C3、C4、C5。这里要注意,torchvision的ResNet输出通道分别是512、1024、2048,跟论文里用的配置一致,如果你的骨干是别的网络,记得先把通道数对齐。
对接的时候有个细节:C3的分辨率是输入的1/8,C4是1/16,C5是1/32。FPT的三个分支都在这三个层级上操作,所以输入图像的尺寸最好是32的整数倍,不然上采样和下采样的尺寸会对不齐,拼接时报错。我一般把输入resize到800x800或者1024x1024,都是32的倍数,省心。
import torch import torch.nn as nn import torchvision class FPTNeck(nn.Module): def __init__(self, in_channels_list, embed_dim=256, num_heads=8): super().__init__() # 统一通道维度的横向连接 self.lateral = nn.ModuleList([ nn.Conv2d(c, embed_dim, 1) for c in in_channels_list ]) # 三个融合分支的注意力模块(简化示意) self.td_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) self.bu_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) self.same_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True) # 融合后的聚合卷积 self.fuse = nn.Conv2d(embed_dim * 3, embed_dim, 1)上面这段是骨架级的示意,实际实现里每个分支还要处理上采样、下采样、reshape这些操作,我把它拆到下面的小节里讲。
4.2 自顶向下分支的完整实现
这个分支的输入是相邻两层特征,高层做query,低层做key和value。实现时先把两层特征都过lateral卷积统一到256通道,然后展平。注意展平前要记录原始的空间尺寸,因为后面要reshape回去。
def top_down_fusion(self, high_feat, low_feat): # high_feat: (B, C, Hh, Wh), low_feat: (B, C, Hl, Wl) B, C, Hh, Wh = high_feat.shape _, _, Hl, Wl = low_feat.shape # 展平为序列 q = high_feat.flatten(2).transpose(1, 2) # (B, Hh*Wh, C) kv = low_feat.flatten(2).transpose(1, 2) # (B, Hl*Wl, C) # 交叉注意力 attn_out, _ = self.td_attn(q, kv, kv) # reshape回空间维度 attn_out = attn_out.transpose(1, 2).reshape(B, C, Hh, Wh) # 上采样到低层分辨率后融合 attn_out = F.interpolate(attn_out, size=(Hl, Wl), mode='bilinear', align_corners=False) return attn_out + low_feat这段代码里有个容易忽略的点:F.interpolate的align_corners参数。设成False是PyTorch的默认推荐值,但在特征图对齐要求高的场景(比如要做像素级分割),设成True可能更合适。我两种都试过,检测任务上差别不大,分割任务上align_corners=True的边缘对齐略好一点,但差异很小,按任务习惯选就行。
4.3 自底向上分支与同层分支的落地
自底向上分支跟上面几乎对称,只是query和key/value的角色互换,聚合完是下采样而不是上采样。下采样我用的是步长2的3x3卷积,比直接池化多了一点可学习性。
def bottom_up_fusion(self, low_feat, high_feat): B, C, Hl, Wl = low_feat.shape _, _, Hh, Wh = high_feat.shape q = low_feat.flatten(2).transpose(1, 2) kv = high_feat.flatten(2).transpose(1, 2) attn_out, _ = self.bu_attn(q, kv, kv) attn_out = attn_out.transpose(1, 2).reshape(B, C, Hl, Wl) # 用步长卷积下采样到高层分辨率 attn_out = self.down_conv(attn_out) # stride=2, 3x3 return attn_out + high_feat同层分支就是标准自注意力,但正如前面说的,全分辨率做会爆显存,所以我加了一个可选的池化预处理。如果特征图边长超过64,先池化到一半再做注意力,算完再插值回去。这个阈值是我根据显存实测定的,你可以根据自己的卡调整。
三个分支的输出尺寸要统一到同一层级才能融合。我的做法是:对每个层级i,自顶向下分支的输出对齐到层级i的分辨率,自底向上分支的输出也对齐到层级i,同层分支本来就是层级i的分辨率,三者拼接后过1x1卷积降维。这样每个层级都得到一个融合了三个方向信息的特征图,再送进检测头。
4.4 训练配置与收敛性观察
训练这块我用的配置比较常规:SGD优化器,初始学习率0.01,momentum 0.9,weight decay 1e-4,batch size 4(受显存限制),训练12个epoch,在第8和第11个epoch各降一次学习率,降为原来的0.1倍。数据增强用了随机翻转、随机缩放和多尺度训练。
收敛性上有个现象值得说:FPT的loss曲线在前两个epoch下降比标准FPN慢一些,我一开始以为是实现有问题,后来对比了注意力权重的分布才明白,是因为注意力模块的初始化让融合权重接近均匀分布,网络需要几个epoch才能学到有区分度的跨尺度对应关系。从第3个epoch开始loss下降速度就追上来了,最终收敛后的精度比标准FPN高。所以如果你训到第2个epoch发现loss还比较高,别急着怀疑代码,再等等。
注意:注意力模块的初始化建议用小方差的正态分布(比如std=0.02),别用默认初始化。我试过默认初始化,前几个epoch的梯度范数波动很大,训练不太稳。
5. 常见问题与排查技巧实录
5.1 显存爆炸的三种典型场景
复现FPT最常遇到的问题就是显存不够,我踩过至少三种不同的爆显存场景,这里逐个说下排查思路。
第一种是同层注意力在P3上全分辨率计算。前面提过,P3序列长度能到一万,注意力矩阵是一万乘一万,显存直接起飞。排查方法很简单,在注意力计算前打印一下序列长度,超过4096就要警惕。解决办法是池化下采样或者窗口化,我一般先用池化快速验证。
第二种是三个分支并行计算时的中间激活值累积。三个分支各自产生一批中间张量,如果不在分支内部及时释放,显存峰值会是单分支的三倍左右。我的做法是在每个分支函数内部算完就返回结果,中间变量靠Python的引用计数自动回收,同时用torch.cuda.empty_cache()在训练循环里定期清理(虽然官方说没必要,但实测在长训练里能缓解碎片化)。
第三种是batch size和输入尺寸的乘积超限。这个最直观,但容易被忽略的是多尺度训练时,最大尺度的那个batch可能刚好卡在显存边缘,训练到一半随机切到大尺度就OOM。我的经验是把最大尺度的batch单独降一档,或者用梯度累积模拟大batch。
5.2 注意力权重全均匀或不收敛的排查
另一个高频问题是注意力权重学不出区分度,可视化出来一片均匀,等于没融合。我遇到过两次,原因不同。
第一次是位置编码没加或者加错了维度。同层注意力对位置信息依赖很强,位置编码缺失时注意力退化成近似平均池化。检查方法是把注意力矩阵打印出来看对角线附近有没有聚集,如果没有,大概率是位置编码的问题。
第二次是学习率太大导致注意力参数震荡。注意力模块的参数量不大,对学习率比较敏感,用骨干网络同款学习率容易训飞。我的处理是给注意力模块单独设一个更小的学习率(比如主学习率的0.1倍),用参数组的方式配置优化器,训下来稳定很多。
还有一种情况是query和key的特征分布差异过大。跨层注意力里,高层和低层的特征统计特性本来就不一样,如果lateral卷积的初始化不好,query和key的点积会落在饱和区,softmax之后接近均匀。解决办法是在lateral卷积后加一个LayerNorm,把特征分布拉齐,这个改动很小但对注意力质量提升明显。
5.3 精度不升反降的可能原因
有时候复现完发现精度还不如标准FPN,别慌,按下面这个表逐项排查,大部分情况能定位到原因。
| 现象 | 可能原因 | 排查与解决 |
|---|---|---|
| 训练loss正常但验证mAP低 | 过拟合,注意力参数量相对数据量偏大 | 加dropout或weight decay,减少注意力头数 |
| 小目标精度下降明显 | 自底向上分支的下采样丢失细节 | 下采样改用步长卷积,或保留高分辨率特征 |
| 大目标精度下降 | 同层注意力过度平滑了空间细节 | 降低同层分支权重,或去掉同层分支 |
| 训练前期loss震荡 | 注意力初始化不当或学习率过大 | 小方差初始化,注意力模块单独降学习率 |
| 推理速度慢到不可接受 | 三分支并行计算开销大 | 去掉同层分支,或对注意力做窗口化 |
我印象最深的一次是去掉同层分支后小目标精度反而涨了,后来分析是因为同层注意力在P3上把邻近小目标的特征平滑到一起了,导致密集小目标互相干扰。所以如果你的场景是小目标密集(比如航拍检测),同层分支要慎用,或者把它的窗口限制得很小。
5.4 推理部署时的算子兼容问题
训练跑通不代表部署顺利。FPT里的多头注意力在导出到某些推理引擎时可能不被支持,或者被拆成一堆小算子导致速度很慢。我的处理方式是导出前把注意力模块里的reshape和transpose操作尽量合并,减少算子数量。另外,如果目标平台对动态shape支持不好,要把序列长度固定下来,避免运行时动态计算。
实测下来,在GPU上FPT的推理耗时比标准FPN多大概30%到40%,主要开销在跨层注意力的矩阵乘上。如果对延迟敏感,可以考虑只在P4和P5之间做跨层融合,P3保持标准FPN的处理,这样能砍掉一大半计算量,精度损失在可接受范围内。
6. 我个人的几点实操体会
这个结构我从看到论文到跑通复现大概花了一周多,中间踩的坑基本都写在上面了。回过头看,最值得分享的一条经验是:别一上来就三个分支全开。我最初就是三个分支一起上,结果显存爆了、训练不稳、精度还不升,一度怀疑论文有问题。后来老老实实从只开自顶向下分支开始,确认这条路径有效果,再加自底向上,最后加同层,每加一个分支做一次消融对比,才把问题定位清楚。这种增量式的复现策略,比一次性全量实现要稳得多。
另外一点是关于“借鉴Transformer”这件事的看法。FPT把注意力搬进特征金字塔,本质上是用一种更灵活的方式替代了固定的融合规则。但注意力不是免费的,它的计算开销和显存占用都是实打实的。所以在决定用不用它之前,先想清楚你的场景里跨尺度信息融合到底是不是瓶颈。如果你的目标尺度分布很集中,标准FPN可能就够了,硬上FPT属于杀鸡用牛刀。反过来,如果你的场景里尺度跨度大、小目标多、语义和细节的对齐要求高,那FPT这套多方向融合的思路确实能带来实打实的收益,值得花时间落地。
最后分享一个小技巧:调试注意力模块时,把注意力权重图可视化出来看,比盯着loss曲线有用得多。我一般会挑几张验证集图片,把每个头的注意力热力图存下来,直观地看它到底在关注哪些区域。很多时候精度不升的原因,看一眼注意力图就明白了——要么全图均匀,要么只关注自己,要么关注到了背景上。这种可视化排查的效率,比盲调超参高一个数量级。