083、ASFF 自适应空间特征融合：Level 0/1/2 自学习融合权重的 Softmax 实现

083、ASFF 自适应空间特征融合Level 0/1/2 自学习融合权重的 Softmax 实现从一次诡异的mAP波动说起去年秋天调YOLOv5的Neck结构跑COCO验证集时发现一个怪现象换了BiFPN之后小目标AP涨了2个点大目标AP却掉了1.5个点。当时盯着TensorBoard的曲线看了半小时心想这玩意儿怎么还挑食呢后来翻FPN论文的消融实验发现不同尺度的特征图在空间位置上存在“语义冲突”——高层语义特征在背景区域激活强低层纹理特征在目标边缘激活强简单相加等于让两个观点不同的人强行握手结果两边都不舒服。ASFFAdaptively Spatial Feature Fusion就是来解决这个问题的。它让网络自己学每个空间位置该信哪个尺度的特征而不是一刀切地加权求和。今天咱们就手撕ASFF的PyTorch实现重点看Level 0/1/2三个尺度的自适应权重是怎么通过Softmax算出来的。ASFF的核心思想别让特征图打架先理解一下ASFF在干什么。假设你有三个尺度的特征图Level 0大尺度比如P3、Level 1中尺度P4、Level 2小尺度P5。传统FPN直接把它们加起来output l0 l1 l2。ASFF改成output alpha * l0 beta * l1 gamma * l2其中alpha、beta、gamma是每个空间位置都不同的权重且满足alpha beta gamma 1。这个权重怎么来的不是拍脑袋定的而是从特征图本身学出来的。具体做法是对每个尺度的特征图分别过一层1x1卷积生成一个单通道的权重图然后把三个权重图拼起来过Softmax得到每个位置三个尺度的归一化权重。代码实现从Level 0/1/2到融合输出先看ASFF的整体结构。假设输入是三个尺度的特征图尺寸分别为Level 0 (B, C, H, W)、Level 1 (B, C, H/2, W/2)、Level 2 (B, C, H/4, W/4)。注意这三个特征图的通道数C必须一致空间尺寸不同。所以第一步要做的是把Level 1和Level 2上采样到Level 0的尺寸。importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassASFF(nn.Module):def__init__(self,level,channels,reduction_ratio16):super().__init__()self.levellevel# 当前融合的目标尺度0/1/2# 每个尺度生成权重图的1x1卷积# 这里踩过坑权重图通道数设为1不是channelsself.weight_convsnn.ModuleList([nn.Conv2d(channels,1,kernel_size1,biasFalse)for_inrange(3)])# 如果是Level 0需要上采样Level 1和Level 2# 如果是Level 1需要下采样Level 0上采样Level 2# 如果是Level 2需要下采样Level 0和Level 1# 别这样写把三个尺度的上采样/下采样都写死后面改网络结构时容易崩# 建议用interpolate动态调整defforward(self,x0,x1,x2):# x0: Level 0, 尺寸最大# x1: Level 1, 尺寸中等# x2: Level 2, 尺寸最小# 获取目标尺寸以当前level的尺寸为准target_h,target_wx0.shape[2:]ifself.level0else\(x1.shape[2:]ifself.level1elsex2.shape[2:])# 将三个尺度的特征图resize到目标尺寸# 这里用F.interpolate别用nn.Upsample后者是固定尺寸的x0_resizedF.interpolate(x0,size(target_h,target_w),modebilinear,align_cornersFalse)x1_resizedF.interpolate(x1,size(target_h,target_w),modebilinear,align_cornersFalse)x2_resizedF.interpolate(x2,size(target_h,target_w),modebilinear,align_cornersFalse)# 生成三个尺度的权重图weight0self.weight_convs[0](x0_resized)# (B, 1, H, W)weight1self.weight_convs[1](x1_resized)weight2self.weight_convs[2](x2_resized)# 拼起来过Softmax得到归一化权重# 别这样写torch.cat([weight0, weight1, weight2], dim1) 然后手动算softmax# 直接用一个catsoftmax搞定weightstorch.cat([weight0,weight1,weight2],dim1)# (B, 3, H, W)weightsF.softmax(weights,dim1)# 在通道维度上归一化# 加权融合# 这里踩过坑weights的shape是(B,3,H,W)需要拆成三个单通道w0weights[:,0:1,:,:]# 保持维度别用squeezew1weights[:,1:2,:,:]w2weights[:,2:3,:,:]outputw0*x0_resizedw1*x1_resizedw2*x2_resizedreturnoutput这段代码看起来简单但有几个细节值得注意。首先是weight_convs的输入输出通道数输入是channels输出是1。为什么是1因为我们要的是每个空间位置的一个标量权重不是特征向量。如果你输出channels个通道那Softmax之后每个位置会有channels个权重每个权重对应一个通道这就变成了通道注意力不是空间注意力了。别这样写除非你想做通道级别的自适应融合。其次是resize操作。ASFF要求三个尺度的特征图在融合前尺寸一致所以需要上采样或下采样。这里统一用F.interpolate模式选bilinearalign_cornersFalse。为什么不用最近邻因为特征图是连续的空间信号双线性插值更平滑。为什么align_cornersFalse这是PyTorch的默认行为对齐像素中心避免边缘偏移。如果你用align_cornersTrue上采样后的特征图会有0.5像素的偏移对于小目标检测来说影响很大。多尺度融合的完整实现三个ASFF模块实际使用时我们需要对每个尺度都做一次ASFF。比如在YOLOv5的Neck中P3、P4、P5三个特征图分别对应Level 0、1、2。我们需要三个ASFF模块每个模块输出一个融合后的特征图尺寸与当前Level一致。classASFFNeck(nn.Module):def__init__(self,channels):super().__init__()# 三个ASFF模块分别对应Level 0/1/2self.asff0ASFF(level0,channelschannels)self.asff1ASFF(level1,channelschannels)self.asff2ASFF(level2,channelschannels)defforward(self,features):# features: [P3, P4, P5] 对应 Level 0/1/2p3,p4,p5features# 每个Level都做一次自适应融合out0self.asff0(p3,p4,p5)# 输出尺寸与P3一致out1self.asff1(p3,p4,p5)# 输出尺寸与P4一致out2self.asff2(p3,p4,p5)# 输出尺寸与P5一致return[out0,out1,out2]这里有个容易忽略的点三个ASFF模块的权重是独立学习的。也就是说Level 0的融合权重只关心如何把P3、P4、P5融合成P3尺寸的特征图Level 1的权重只关心如何融合成P4尺寸。它们之间没有共享参数因为不同尺度的语义信息不同需要的融合策略也不同。权重可视化看看网络学到了什么训练完成后可以把权重图可视化出来。你会发现一个有趣的现象在目标区域Level 0大尺度的权重通常较高因为大尺度特征图保留了更多细节在背景区域Level 2小尺度的权重较高因为高层语义特征对背景的响应更稳定。这正好印证了ASFF的设计初衷——让网络自己决定每个位置该信谁。可视化代码很简单# 假设已经跑了一次forward拿到了weightsweightsweights.detach().cpu().numpy()# (B, 3, H, W)# 取batch中第一张图w0weights[0,0,:,:]# Level 0的权重图w1weights[0,1,:,:]# Level 1的权重图w2weights[0,2,:,:]# Level 2的权重图# 用matplotlib或opencv显示即可注意权重图是浮点数范围在0到1之间可以直接用热力图显示。如果你发现某个尺度的权重几乎全是0或全是1说明网络没有学到有效的融合策略可能是训练不充分或者特征图本身已经足够好了不需要自适应融合。踩坑记录那些年我掉过的坑梯度消失问题如果权重图的初始值太大或太小Softmax之后会接近one-hot导致梯度消失。解决办法是在weight_convs的初始化中把权重设小一点比如nn.init.normal_(conv.weight, mean0, std0.01)。内存爆炸ASFF需要同时保留三个尺度的特征图和三个权重图显存占用是普通FPN的2倍左右。如果你的GPU显存只有8G建议把reduction_ratio设大一点或者只在Neck的最后几层用ASFF。训练不稳定刚开始训练时权重图可能剧烈变化导致loss震荡。可以在前几个epoch固定权重为均匀分布即每个尺度权重都是1/3等网络稳定后再放开。实现方式是在forward里加一个self.training判断。与BN的兼容性ASFF的权重图是1x1卷积生成的后面没有BN层。如果你在weight_convs后面加BN反而会破坏权重的尺度因为BN会强制输出均值为0、方差为1而Softmax需要输入是任意实数。别这样写直接不加BN。个人经验什么时候该用ASFFASFF不是万能的。如果你的数据集目标尺度分布很均匀比如全是中等大小的物体那普通FPN就够用了。ASFF的优势在于处理多尺度目标尤其是小目标和超大目标共存的情况。比如自动驾驶场景近处的行人大目标和远处的交通标志小目标同时出现ASFF能显著提升小目标的召回率。另外ASFF的计算开销主要在三个1x1卷积和三次resize上。如果你用YOLOv5s这种轻量模型ASFF可能会让推理速度下降10%-15%。建议在YOLOv5m或更大的模型上使用性价比更高。最后说一句ASFF的权重可视化是个很好的调试工具。如果你发现某个尺度的权重始终很低说明这个尺度的特征图对当前任务贡献不大可以考虑直接去掉或者用更轻量的融合方式。别盲目堆模块理解你的数据比理解代码更重要。