CVPR | CNN融合注意力机制，芜湖起飞！

**标题：**On the Integration of Self-Attention and Convolution
**论文链接：**https://arxiv.org/pdf/2111.14556
**代码链接：**https://github.com/LeapLabTHU/ACmix

创新点

1. 揭示卷积和自注意力的内在联系

文章通过重新分解卷积和自注意力模块的操作，发现它们在第一阶段（特征投影）都依赖于 1×1 卷积操作，并且这一阶段占据了大部分的计算复杂度（与通道数的平方成正比）。这一发现为整合两种模块提供了理论基础。

2. 提出 ACmix 模型

基于上述发现，作者提出了 ACmix 模型，它通过共享 1×1 卷积操作来同时实现卷积和自注意力的功能。具体来说：
**第一阶段：**输入特征通过 1×1 卷积投影，生成中间特征。
**第二阶段：**这些中间特征分别用于卷积路径（通过移位和聚合操作）和自注意力路径（计算注意力权重并聚合值）。最终，两条路径的输出通过可学习的权重加权求和，得到最终输出。

3. 改进的移位和聚合操作

文章还提出了一种改进的移位操作，通过使用固定卷积核的分组卷积来替代传统的张量移位操作。这种方法不仅提高了计算效率，还允许卷积核的可学习性，进一步增强了模型的灵活性。

4. 适应性路径权重

ACmix 引入了两个可学习的标量参数（α 和 β），用于动态调整卷积路径和自注意力路径的权重。这种设计不仅提高了模型的灵活性，还允许模型在不同深度上自适应地选择更适合的特征提取方式。实验表明，这种设计在模型的不同阶段表现出不同的偏好，例如在早期阶段更倾向于卷积，在后期阶段更倾向于自注意力。

整体结构

第一阶段：特征投影

在第一阶段，输入特征通过三个1×1卷积进行投影，分别生成查询（query）、键（key）和值（value）特征映射。这些特征映射随后被重塑为N块，形成一个包含3×N特征映射的中间特征集。

第二阶段：特征聚合

在第二阶段，中间特征集被分为两个路径进行处理：

**自注意力路径：**将中间特征集分为N组，每组包含三个特征映射（分别对应查询、键和值）。这些特征映射按照传统的多头自注意力机制进行处理，计算注意力权重并聚合值。
**卷积路径：**通过轻量级的全连接层生成k²个特征映射（k为卷积核大小）。这些特征映射通过移位和聚合操作，以类似传统卷积的方式处理输入特征，从局部感受野收集信息。

输出整合

最后，自注意力路径和卷积路径的输出通过两个可学习的标量参数（α和β）加权求和，得到最终的输出。

改进的移位和聚合操作

为了提高计算效率，ACmix模型采用了改进的移位操作，通过固定卷积核的分组卷积来替代传统的张量移位操作。这种方法不仅提高了计算效率，还允许卷积核的可学习性，进一步增强了模型的灵活性。

模型的灵活性和泛化能力

ACmix模型不仅适用于标准的自注意力机制，还可以与各种变体（如Patchwise Attention、Window Attention和Global Attention）结合使用。这种设计使得ACmix能够适应不同的任务需求，具有广泛的适用性。

消融实验

1. 结合两个路径的输出

消融实验探索了卷积和自注意力输出的不同组合方式对模型性能的影响。实验结果表明：

**卷积和自注意力的组合优于单一路径：**使用卷积和自注意力模块的组合始终优于仅使用单一路径（如仅卷积或仅自注意力）的模型。
**可学习参数的灵活性：**通过引入可学习的参数（如α和β）来动态调整卷积和自注意力路径的权重，ACmix能够根据网络中不同位置的需求自适应地调整路径强度，从而获得更高的灵活性和性能。

2. 组卷积核的选择

实验还对组卷积核的设计进行了验证，结果表明：

**用组卷积替代张量位移：**通过使用组卷积替代传统的张量位移操作，显著提高了模型的推理速度。
**可学习卷积核和初始化：**使用可学习的卷积核并结合精心设计的初始化方法，进一步增强了模型的灵活性，并有助于提升最终性能。

3. 不同路径的偏好

ACmix模型引入了两个可学习标量α和β，用于动态调整卷积和自注意力路径的权重。通过平行实验，观察到以下趋势：

**早期阶段偏好卷积：**在Transformer模型的早期阶段，卷积作为特征提取器表现更好。
**中间阶段混合使用：**在网络的中间阶段，模型倾向于混合使用两种路径，并逐渐增加对卷积的偏好。
**后期阶段偏好自注意力：**在网络的最后阶段，自注意力表现优于卷积。

4. 对模型性能的影响

这些消融实验结果表明，ACmix模型通过合理结合卷积和自注意力的优势，并优化计算路径，不仅在多个视觉任务上取得了显著的性能提升，还保持了较高的计算效率

ACmix模块的作用

1. 融合卷积和自注意力的优势

ACmix模块通过结合卷积的局部特征提取能力和自注意力的全局感知能力，实现了一种高效的特征融合策略。这种设计使得模型能够同时利用卷积的局部感受野特性和自注意力的灵活性。

2. 优化计算路径

ACmix通过优化计算路径和减少重复计算，提高了整体模块的计算效率。具体来说，它通过1×1卷积对输入特征图进行投影，生成中间特征，然后根据不同的范式（卷积和自注意力）分别重用和聚合这些中间特征。这种设计不仅减少了计算开销，还提升了模型性能。

3. 改进的位移与求和操作

在卷积路径中，ACmix采用深度可分离卷积（depthwise convolution）来替代低效的张量位移操作，从而提高了实际推理效率。

4. 动态调整路径权重

ACmix引入了两个可学习的标量参数（α和β），用于动态调整卷积和自注意力路径的权重。这种设计使得模型能够根据网络中不同位置的需求自适应地调整路径强度，从而获得更高的灵活性。

5. 广泛的应用潜力

ACmix在多个视觉任务（如图像分类、语义分割和目标检测）上均显示出优于单一机制（仅卷积或仅自注意力）的性能，展示了其广泛的应用潜力。

6. 实验验证

实验结果表明，ACmix在保持较低计算开销的同时，能够显著提升模型的性能。例如，在ImageNet分类任务中，ACmix模型在相同的FLOPs或参数数量下表现出色，并且在与竞争对手的基准比较中取得了持续的改进。此外，ACmix在ADE20K语义分割任务和COCO目标检测任务中也显示出明显的改进

代码实现

import torch
import torch.nn as nn


def position(H, W, is_cuda=True):
    if is_cuda:
        loc_w = torch.linspace(-1.0, 1.0, W).cuda().unsqueeze(0).repeat(H, 1)
        loc_h = torch.linspace(-1.0, 1.0, H).cuda().unsqueeze(1).repeat(1, W)
    else:
        loc_w = torch.linspace(-1.0, 1.0, W).unsqueeze(0).repeat(H, 1)
        loc_h = torch.linspace(-1.0, 1.0, H).unsqueeze(1).repeat(1, W)
    loc = torch.cat([loc_w.unsqueeze(0), loc_h.unsqueeze(0)], 0).unsqueeze(0)
    return loc


def stride(x, stride):
    b, c, h, w = x.shape
    return x[:, :, ::stride, ::stride]


def init_rate_half(tensor):
    if tensor is not None:
        tensor.data.fill_(0.5)


def init_rate_0(tensor):
    if tensor is not None:
        tensor.data.fill_(0.)


class ACmix(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_att=7, head=4, kernel_conv=3, stride=1, dilation=1):
        super(ACmix, self).__init__()
        self.in_planes = in_planes
        self.out_planes = out_planes
        self.head = head
        self.kernel_att = kernel_att
        self.kernel_conv = kernel_conv
        self.stride = stride
        self.dilation = dilation
        self.rate1 = torch.nn.Parameter(torch.Tensor(1))
        self.rate2 = torch.nn.Parameter(torch.Tensor(1))
        self.head_dim = self.out_planes // self.head

        self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=1)
        self.conv2 = nn.Conv2d(in_planes, out_planes, kernel_size=1)
        self.conv3 = nn.Conv2d(in_planes, out_planes, kernel_size=1)
        self.conv_p = nn.Conv2d(2, self.head_dim, kernel_size=1)

        self.padding_att = (self.dilation * (self.kernel_att - 1) + 1) // 2
        self.pad_att = torch.nn.ReflectionPad2d(self.padding_att)
        self.unfold = nn.Unfold(kernel_size=self.kernel_att, padding=0, stride=self.stride)
        self.softmax = torch.nn.Softmax(dim=1)

        self.fc = nn.Conv2d(3 * self.head, self.kernel_conv * self.kernel_conv, kernel_size=1, bias=False)
        self.dep_conv = nn.Conv2d(self.kernel_conv * self.kernel_conv * self.head_dim, out_planes,
                                  kernel_size=self.kernel_conv, bias=True, groups=self.head_dim, padding=1,
                                  stride=stride)

        self.reset_parameters()

    def reset_parameters(self):
        init_rate_half(self.rate1)
        init_rate_half(self.rate2)
        kernel = torch.zeros(self.kernel_conv * self.kernel_conv, self.kernel_conv, self.kernel_conv)
        for i in range(self.kernel_conv * self.kernel_conv):
            kernel[i, i // self.kernel_conv, i % self.kernel_conv] = 1.
        kernel = kernel.squeeze(0).repeat(self.out_planes, 1, 1, 1)
        self.dep_conv.weight = nn.Parameter(data=kernel, requires_grad=True)
        self.dep_conv.bias = init_rate_0(self.dep_conv.bias)

    def forward(self, x):
        q, k, v = self.conv1(x), self.conv2(x), self.conv3(x)
        scaling = float(self.head_dim) ** -0.5
        b, c, h, w = q.shape
        h_out, w_out = h // self.stride, w // self.stride

        # ### att
        # ## positional encoding
        pe = self.conv_p(position(h, w, x.is_cuda))

        q_att = q.view(b * self.head, self.head_dim, h, w) * scaling
        k_att = k.view(b * self.head, self.head_dim, h, w)
        v_att = v.view(b * self.head, self.head_dim, h, w)

        if self.stride > 1:
            q_att = stride(q_att, self.stride)
            q_pe = stride(pe, self.stride)
        else:
            q_pe = pe

        unfold_k = self.unfold(self.pad_att(k_att)).view(b * self.head, self.head_dim,
                                                         self.kernel_att * self.kernel_att, h_out,
                                                         w_out) # b*head, head_dim, k_att^2, h_out, w_out
        unfold_rpe = self.unfold(self.pad_att(pe)).view(1, self.head_dim, self.kernel_att * self.kernel_att, h_out,
                                                        w_out) # 1, head_dim, k_att^2, h_out, w_out

        att = (q_att.unsqueeze(2) * (unfold_k + q_pe.unsqueeze(2) - unfold_rpe)).sum(
            1) # (b*head, head_dim, 1, h_out, w_out) * (b*head, head_dim, k_att^2, h_out, w_out) -> (b*head, k_att^2, h_out, w_out)
        att = self.softmax(att)

        out_att = self.unfold(self.pad_att(v_att)).view(b * self.head, self.head_dim, self.kernel_att * self.kernel_att,
                                                        h_out, w_out)
        out_att = (att.unsqueeze(1) * out_att).sum(2).view(b, self.out_planes, h_out, w_out)

        ## conv
        f_all = self.fc(torch.cat(
            [q.view(b, self.head, self.head_dim, h * w), k.view(b, self.head, self.head_dim, h * w),
             v.view(b, self.head, self.head_dim, h * w)], 1))
        f_conv = f_all.permute(0, 2, 1, 3).reshape(x.shape[0], -1, x.shape[-2], x.shape[-1])

        out_conv = self.dep_conv(f_conv)

        return self.rate1 * out_att + self.rate2 * out_conv


#输入 B C H W, 输出 B C H W
if __name__ == '__main__':
    block = ACmix(in_planes=64, out_planes=64)
    input = torch.rand(3, 64, 32, 32)
    output = block(input)
    print(input.size(), output.size())