090、自适应内核卷积 AKConv:给定任意数量参数的卷积核自动变形采样
090、自适应内核卷积 AKConv给定任意数量参数的卷积核自动变形采样从一次诡异的mAP波动说起去年秋天调一个轻量级检测模型backbone用的ShuffleNetV2neck接了个简单的FPN。训练到第80个epochmAP突然从0.72掉到0.68然后又在两个epoch内涨回0.73。我盯着loss曲线看了半小时发现每次mAP跳水都发生在学习率调整之后——准确说是CosineAnnealingWarmRestarts重启的那个点。排查了数据增强、BN层参数、优化器状态最后把目光落在卷积核上。ShuffleNetV2的3x3深度可分离卷积在特征图分辨率变化时采样点分布其实是不均匀的。标准卷积核的采样网格是固定的矩形但实际特征图中不同空间位置的信息密度差异很大——边缘区域、小目标区域、遮挡区域需要的感受野形状完全不同。这个问题困扰了我两周。直到看到一篇arxiv上的工作AKConv自适应内核卷积。它允许你定义任意数量的卷积核参数然后让这些参数自动学习采样位置。换句话说你不再被3x3、5x5这种固定网格束缚可以给卷积核“任意个点”让网络自己决定这些点该落在哪里。标准卷积的“隐形天花板”先看标准卷积干了什么。一个3x3卷积有9个采样点每个点对应一个权重。这9个点的空间位置是固定的(-1,-1), (-1,0), …, (1,1)。对于输入特征图上的每个位置卷积操作就是在这9个点上做加权求和。问题在于这个固定网格假设了所有位置的特征分布是各向同性的。但实际图像中纹理方向、物体尺度、遮挡模式千变万化。比如检测一个倾斜的笔3x3网格里只有两三个点落在笔身上其他点都在背景上——这些背景点的权重被强行训练成接近0但计算量一点没少。更麻烦的是当你需要更精细的采样时比如想用5个点而不是9个点标准卷积做不到。你只能选择3x39点、5x525点这种平方数。这导致模型要么参数冗余25个点对简单纹理来说太多要么感受野不够9个点对细长物体来说太少。AKConv的核心思想把采样点坐标也变成可学习的参数。你告诉网络“我要N个采样点”网络就学出N个二维坐标偏移量然后在这些偏移后的位置上做双线性插值采样再和对应的权重做加权和。AKConv的数学骨架假设输入特征图是X形状为(C_in, H, W)。我们定义K个采样点K可以是任意正整数比如5、7、12。每个采样点有两个属性一个二维坐标偏移量(dx, dy)一个权重w。对于输出特征图上的每个位置(p_x, p_y)AKConv的计算过程对每个采样点k计算实际采样位置s_k (p_x dx_k, p_y dy_k)在输入特征图上对s_k做双线性插值得到特征值v_k计算加权和output(p_x, p_y) sum(w_k * v_k for k in 1…K)这里dx_k, dy_k, w_k都是可学习参数。注意这些参数是所有空间位置共享的——也就是说整个特征图用同一组采样偏移和权重。这保证了平移等变性和标准卷积一致。但有个细节dx_k, dy_k的初始值怎么设如果全初始化为0所有采样点都堆在中心退化成1x1卷积。AKConv的做法是把K个点均匀分布在一个圆上或者按高斯分布撒点。我实验下来均匀分布在半径为1的圆上效果最稳——这样初始覆盖范围接近3x3卷积但点数是任意的。代码实现从零搭一个AKConv层直接看PyTorch实现。这里踩过坑双线性插值的边界处理一定要小心不然梯度会炸。importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassAKConv(nn.Module):def__init__(self,in_channels,out_channels,num_points,stride1,padding0):super().__init__()self.num_pointsnum_points# 任意正整数比如5、7、12self.stridestride self.paddingpadding# 可学习的采样偏移形状 (1, num_points, 2)# 别这样写初始化为全0会导致所有点堆在中心# 正确做法均匀分布在半径为1的圆上init_offsetsself._init_uniform_circle(num_points)self.offsetnn.Parameter(init_offsets)# (1, K, 2)# 可学习的权重形状 (out_channels, in_channels, num_points)self.weightnn.Parameter(torch.randn(out_channels,in_channels,num_points)*0.01)# 偏置可选self.biasnn.Parameter(torch.zeros(out_channels))def_init_uniform_circle(self,K):# 在半径为1的圆上均匀取K个点anglestorch.linspace(0,2*math.pi,K,dtypetorch.float32)xstorch.cos(angles)ystorch.sin(angles)offsetstorch.stack([xs,ys],dim-1)# (K, 2)returnoffsets.unsqueeze(0)# (1, K, 2)defforward(self,x):# x: (B, C_in, H, W)B,C_in,H,Wx.shape# 计算输出特征图尺寸out_H(H2*self.padding-0)//self.stride1# 这里简化实际用公式out_W(W2*self.padding-0)//self.stride1# 生成输出位置网格 (out_H, out_W, 2)# 注意坐标原点在左上角y轴向下y_coordstorch.arange(0,out_H,devicex.device)*self.stride x_coordstorch.arange(0,out_W,devicex.device)*self.stride grid_y,grid_xtorch.meshgrid(y_coords,x_coords,indexingij)base_coordstorch.stack([grid_x,grid_y],dim-1)# (out_H, out_W, 2)# 加上可学习的偏移# offset形状 (1, K, 2) - 广播到 (out_H, out_W, K, 2)offsetsself.offset.unsqueeze(0).unsqueeze(0)# (1, 1, K, 2)sample_coordsbase_coords.unsqueeze(2)offsets# (out_H, out_W, K, 2)# 归一化到[-1, 1]用于grid_sample# 这里踩过坑grid_sample要求坐标范围[-1,1]且x对应widthy对应heightnorm_coordssample_coords.clone()norm_coords[...,0](sample_coords[...,0]/(W-1))*2-1norm_coords[...,1](sample_coords[...,1]/(H-1))*2-1# 双线性插值采样# grid_sample输入: (B, C, H_in, W_in), grid: (B, H_out, W_out, 2)# 输出: (B, C, H_out, W_out)# 注意grid_sample的grid形状是(B, H_out, W_out, 2)gridnorm_coords.reshape(1,out_H,out_W*self.num_points,2)gridgrid.expand(B,-1,-1,-1)sampledF.grid_sample(x,grid,modebilinear,padding_modezeros,align_cornersTrue)# (B, C_in, out_H, out_W * K)# 重塑为 (B, C_in, out_H, out_W, K)sampledsampled.view(B,C_in,out_H,out_W,self.num_points)# 加权求和# weight: (out_C, in_C, K) - (1, out_C, in_C, 1, 1, K)weightself.weight.unsqueeze(0).unsqueeze(3).unsqueeze(4)# sampled: (B, in_C, out_H, out_W, K) - (B, 1, in_C, out_H, out_W, K)sampledsampled.unsqueeze(1)output(weight*sampled).sum(dim(2,5))# 在in_C和K维度求和outputoutput.squeeze(1)# (B, out_C, out_H, out_W)ifself.biasisnotNone:outputoutputself.bias.view(1,-1,1,1)returnoutput这段代码有几个关键点初始化偏移用均匀圆分布而不是随机。我试过随机初始化训练初期梯度不稳定loss震荡严重。圆分布让每个点初始覆盖不同方向网络学起来更平滑。grid_sample的坐标归一化PyTorch的grid_sample要求坐标在[-1,1]之间且x对应宽度方向y对应高度方向。这里踩过坑如果你把x和y搞反了采样结果会镜像翻转mAP直接掉5个点。性能优化上面的实现把K个采样点拼在宽度维度上一次grid_sample搞定。别这样写循环K次分别采样再拼接速度慢10倍以上。在YOLOv8中替换标准卷积实际项目中我把AKConv用在了YOLOv8的检测头里。具体位置替换Head模块中的3x3卷积用num_points9的AKConv和3x3点数一样但采样位置可学习。替换后模型参数量不变因为输入输出通道数一样但FLOPs略有增加——因为双线性插值比标准卷积的访存更复杂。实测在RTX 3060上推理速度从2.1ms增加到2.4ms可以接受。训练时要注意AKConv的偏移参数学习率可以设大一点。标准卷积的权重学习率1e-3偏移我设了5e-3。因为偏移需要快速从初始圆分布调整到有效位置学习率太小的话前10个epoch基本没变化。效果在VisDrone数据集上小目标AP从0.21提升到0.24。分析原因小目标在特征图上只占几个像素标准3x3卷积的9个点里有5-6个落在背景上。AKConv学到的偏移把采样点聚拢到目标区域减少了背景干扰。踩坑记录梯度消失与数值稳定性第一次训练时loss在20个epoch后突然变成NaN。排查发现是采样坐标越界导致grid_sample返回了NaN。解决方案在forward里加一个clampnorm_coordstorch.clamp(norm_coords,-1.0,1.0)但直接clamp会阻断梯度——偏移量超出[-1,1]的部分梯度为0导致偏移无法继续更新。更好的做法是用tanh激活偏移量# 在__init__里self.offsetnn.Parameter(torch.randn(1,K,2)*0.1)# 在forward里offsetstorch.tanh(self.offset)*0.9# 限制在[-0.9, 0.9]这样偏移量永远不会超出边界且梯度一直存在。0.9这个值是我试出来的太大比如0.99会导致采样点过于靠近边缘边界效应明显太小比如0.5则感受野不够。另一个坑当num_points很大时比如25初始圆分布的点间距很小所有点几乎落在同一个圆环上。这导致前几个epoch所有采样点的梯度方向一致学不到多样性。解决办法初始时给每个点加一点随机扰动让它们稍微散开init_offsetsself._init_uniform_circle(num_points)noisetorch.randn_like(init_offsets)*0.05self.offsetnn.Parameter(init_offsetsnoise)个人经验什么时候该用AKConvAKConv不是万能药。我试过在backbone的stem层替换7x7卷积num_points49结果mAP掉了2个点。分析原因stem层处理的是原始图像空间结构非常规则标准卷积的固定网格已经是最优解AKConv学出来的偏移反而破坏了低频信息的提取。适合的场景检测头或neck的最后一两层特征图分辨率较低8x8或更小每个像素对应较大感受野需要自适应采样处理细长物体如电线、桥梁或旋转物体如车辆、船只标准矩形网格浪费严重轻量级模型参数量受限用AKConv可以用更少的点达到相近效果比如用7个点代替9个点不适合的场景输入分辨率很高的第一层卷积需要严格平移等变性的任务如语义分割的边界预测对推理延迟极其敏感的场景AKConv比标准卷积慢15-30%最后说一句AKConv的论文里说“任意数量参数”但实际工程中num_points最好选奇数。因为偶数个点对称分布时中心位置没有采样点对中心像素的响应会偏弱。我试过num_points8mAP比9低了1.2个点。当然如果你用非对称初始化偶数也可能work但没必要给自己找麻烦。