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

1次阅读

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

传统卷积神经网络（CNN）在计算机视觉任务中表现出色，但在处理长距离依赖关系时存在明显局限性。标准的 3 ×3 卷积核感受野有限，难以捕捉图像中远距离像素间的关联。虽然可以通过堆叠更多卷积层来扩大感受野，但这会导致参数数量急剧增加，计算成本上升。

c3k2 模块是一种改进的卷积结构，通过特定的卷积核设计和排列方式，能够在保持较小参数量的同时提供更大的有效感受野。然而，单纯依靠卷积操作仍然无法充分建模全局依赖关系。自注意力机制的引入，正好弥补了这一缺陷，使模型能够直接关注所有位置的特征，并根据它们的相关性进行加权聚合。

让我们先来比较几种常见卷积结构的计算特性：

普通 3 ×3 卷积：
参数量：9×C_in×C_out（假设输入输出通道数分别为 C_in 和 C_out）
计算量：9×C_in×C_out×H×W（H,W 为特征图高宽）
空洞卷积（Dilation=2）：
参数量：同普通 3 ×3 卷积
计算量：同普通 3 ×3 卷积
有效感受野：5×5
c3k2 模块：
参数量：约普通 3 ×3 卷积的 2 /3
计算量：约普通 3 ×3 卷积的 2 /3
有效感受野：通过特定排列可达到 5 ×5 效果

加入自注意力后，虽然会增加一些计算开销，但能显著提升模型对全局信息的建模能力。关键在于如何平衡这两者的比例。

下面是用 PyTorch 实现带自注意力的 c3k2 模块的完整代码：

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

class C3k2WithSelfAttention(nn.Module):
    def __init__(self, in_channels, out_channels, heads=4):
        super().__init__()
        # c3k2 部分
        self.conv1 = nn.Conv2d(in_channels, out_channels//2, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(in_channels, out_channels//2, kernel_size=3, padding=1)

        # 自注意力部分
        self.heads = heads
        self.scale = (out_channels // heads) ** -0.5
        self.to_qkv = nn.Conv2d(out_channels, out_channels * 3, kernel_size=1)
        self.to_out = nn.Conv2d(out_channels, out_channels, kernel_size=1)

    def forward(self, x):
        # c3k2 分支
        x1 = self.conv1(x)
        x2 = self.conv2(x)
        x = torch.cat([x1, x2], dim=1)

        # 自注意力分支
        b, c, h, w = x.shape
        qkv = self.to_qkv(x).chunk(3, dim=1)
        q, k, v = map(lambda t: t.view(b, self.heads, -1, h * w), qkv)

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

        out = torch.matmul(v, attn.transpose(-1, -2))
        out = out.view(b, -1, h, w)
        return self.to_out(out) + x  # 残差连接

我们在 ImageNet-1k 的 100 类子集上进行了对比实验，结果如下：