深入解析c3k2模块加自注意力机制：原理、实现与性能优化

1次阅读

没有评论

共计 1927 个字符，预计需要花费 5 分钟才能阅读完成。

传统卷积神经网络 (CNN) 通过局部感受野和权重共享高效提取特征，但在处理长距离依赖时面临显著挑战：

感受野限制：3×3 卷积核需要堆叠多层才能建立远距离关联，导致信息传递效率低下
计算冗余 ：为扩大感受野采用空洞卷积时，会引入网格效应(gridding artifact) 破坏局部连续性
静态权重：卷积核参数固定，无法根据输入内容动态调整特征提取方式

c3k2 模块 (3×3 卷积与 2×2 深度卷积组合) 通过分解卷积操作减少参数量，而自注意力机制能：

直接建模任意位置间的关系
根据输入动态生成注意力权重
避免远距离信息的多层传递损失

假设输入输出通道均为 C，特征图尺寸 H×W，关键指标对比：

类型	参数量	计算量(FLOPs)	感受野
普通 3×3 卷积	9C²	9C²HW	局部
空洞卷积(rate=2)	9C²	9C²HW	5×5
c3k2 模块	5C²	5C²HW	3×3
c3k2+ 自注意力	5C²+2Ck	5C²HW+2CHWk	全局

注：k 表示注意力头的维度，通常取 C /8

import torch
import torch.nn as nn
import torch.nn.functional as F

class C3K2_Attention(nn.Module):
    def __init__(self, in_channels, head_dim=32):
        super().__init__()
        # 分解卷积部分
        self.conv3x3 = nn.Conv2d(in_channels, in_channels, 3, padding=1, groups=in_channels)
        self.conv1x1 = nn.Conv2d(in_channels, in_channels, 1)

        # 自注意力部分
        self.head_dim = head_dim
        self.num_heads = in_channels // head_dim
        self.scale = head_dim ** -0.5

        self.qkv = nn.Conv2d(in_channels, in_channels * 3, 1)
        self.proj = nn.Conv2d(in_channels, in_channels, 1)

    def forward(self, x):
        # 分解卷积路径
        conv_path = self.conv1x1(self.conv3x3(x))

        # 注意力路径
        B, C, H, W = x.shape
        qkv = self.qkv(x).chunk(3, dim=1)
        q, k, v = map(lambda t: t.view(B, self.num_heads, self.head_dim, H * W), qkv)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)

        out = (attn @ v).view(B, C, H, W)
        out = self.proj(out)

        return conv_path + out  # 残差连接

特征图可视化对比显示：