使用 PyTorch 实现注意力增强卷积网络

• 论文中是用 TensorFlow 实现的，所以我用 PyTorch 重新实现了一遍。

更新（2019年5月11日）

• 修复了在使用 relative=True 模式时，无法找到 key_rel_w 和 key_rel_h 作为可学习参数的问题。

• 在 relative = True 模式下，可以看到 key_rel_w 和 key_rel_h 是可学习的参数。而在 relative = False 模式下，则无需担心 shape 参数。
  
  

• 示例：relative=True，步幅为1，shape=32

import torch

from attention_augmented_conv import AugmentedConv

use_cuda = torch.cuda.is_available()
device = torch.device('cuda' if use_cuda else 'cpu')

tmp = torch.randn((16, 3, 32, 32)).to(device)
augmented_conv1 = AugmentedConv(in_channels=3, out_channels=20, kernel_size=3, dk=40, dv=4, Nh=4, relative=True, stride=1, shape=32).to(device)
conv_out1 = augmented_conv1(tmp)
print(conv_out1.shape) # (16, 20, 32, 32)

for name, param in augmented_conv1.named_parameters():
    print('parameter name: ', name)

• 从参数名称中可以看出，确实存在 key_rel_w 和 key_rel_h。

• 示例：relative=True，步幅为2，shape=16

import torch

from attention_augmented_conv import AugmentedConv

use_cuda = torch.cuda.is_available()
device = torch.device('cuda' if use_cuda else 'cpu')

tmp = torch.randn((16, 3, 32, 32)).to(device)
augmented_conv1 = AugmentedConv(in_channels=3, out_channels=20, kernel_size=3, dk=40, dv=4, Nh=4, relative=True, stride=2, shape=16).to(device)
conv_out1 = augmented_conv1(tmp)
print(conv_out1.shape) # (16, 20, 16, 16)

• 需要注意的是，在使用 relative = True 模式时，步幅乘以 shape 的结果应等于输入图像的尺寸。例如，如果输入是 (16, 3, 32, 32)，且步幅为2，则 shape 应设置为16。

更新（2019年5月2日）

• 我为“AugmentedConv”部分添加了填充功能。

• 使用方式与 nn.conv2d 相同。

• 下面也会附上示例。
  
  

• 示例：relative=False，步幅为1

import torch

from attention_augmented_conv import AugmentedConv

use_cuda = torch.cuda.is_available()
device = torch.device('cuda' if use_cuda else 'cpu')

temp_input = torch.randn((16, 3, 32, 32)).to(device)
augmented_conv = AugmentedConv(in_channels=3, out_channels=20, kernel_size=3, dk=40, dv=4, Nh=1, relative=False，stride=1).to(device)
conv_out = augmented_conv(tmp)
print(conv_out.shape) # (16, 20, 32, 32)，(batch_size, out_channels, height, width)

• 示例：relative=False，步幅为2

import torch

from attention_augmented_conv import AugmentedConv

use_cuda = torch.cuda.isAvailable()
device = torch.device('cuda' if use_cuda else 'cpu')

temp_input = torch.randn((16, 3, 32, 32)).to(device)
augmented_conv = AugmentedConv(in_channels=3, out_channels=20，kernel_size=3，dk=40，dv=4，Nh=1，relative=False，stride=2).to(device)
conv_out = augmented_conv(tmp)
print(conv_out.shape) # (16, 20，16，16)，(batch_size，out_channels，height，width)

• 我为参数（dk、dv、Nh）添加了断言检查。

assert self.Nh != 0, "integer division or modulo by zero, Nh >= 1"
assert self.dk % self.Nh == 0, "dk should be divided by Nh. (example: out_channels: 20, dk: 40, Nh: 4)"
assert self.dv % self.Nh == 0, "dv should be divided by Nh. (example: out_channels: 20, dv: 4, Nh: 4)"
assert stride in [1, 2], str(stride) + " Up to 2 strides are allowed."

我发布了两种版本的“注意力增强卷积”

• 论文版本在此处：https://github.com/leaderj1001/Attention-Augmented-Conv2d/blob/master/attention_augmented_conv.py
• AA-Wide-ResNet 版本在此处：https://github.com/leaderj1001/Attention-Augmented-Conv2d/blob/master/AA-Wide-ResNet/attention_augmented_conv.py

参考资料

论文

• 注意力增强卷积网络论文
• 作者：Irwan Bello、Barret Zoph、Ashish Vaswani、Jonathon Shlens
• Quoc V.Le Google Brain

Wide-ResNet

• GitHub 地址
• 谢谢！

方法

输入参数

• 在论文中， 和 是通过以下公式计算得到的。
  
  ,

• 论文中关于参数的实验
  
  

实验

• 论文中提到，他们通过在所有残差块的第一层卷积中加入相对注意力机制来增强 Wide-ResNet-28-10，使用 Nh=8 头，κ=2，υ=0.2，并且每个头的键至少有 20 维。

数据集	模型	准确率	轮数	训练时间
CIFAR-10	Wide-ResNet 28x10（正在进行中）
CIFAR-100	Wide-ResNet 28x10（正在进行中）
CIFAR-100	仅3层卷积（通道数：64、128、192）	61.6%	100	22分钟
CIFAR-100	仅3层注意力增强卷积（通道数：64、128、192）	59.82%	35	2小时23分钟

• 我没有足够的 GPU，因此训练起来非常困难。
• 我只是想看看这种方法（注意力增强卷积层）是否可行，所以会尝试一下 ResNet。
• 上述结果显示，不同方法之间存在很大的时间差异。我会再深入思考这一部分。
• 我注意到 torch.einsum 函数执行速度较慢。链接
• 当我运行该链接中的示例代码时，结果如下：
  
  
• 使用 CUDA
  
  

时间复杂度

• 我比较了“relative = True”和“relative = False”的时间复杂度。
• 我将对比这两种不同取值（relative=True 和 relative=False）的性能。
• 此外，我还会考虑如何降低“relative = True”情况下的时间复杂度。
  

依赖项

• tqdm==4.31.1
• torch==1.0.1
• torchvision==0.2.2

Attention-Augmented-Conv2d 快速上手指南

本指南帮助中国开发者快速在 PyTorch 环境中部署和使用注意力增强卷积网络（Attention-Augmented Convolutional Networks）。

环境准备

系统要求

• 操作系统: Linux / macOS / Windows
• Python: 建议 Python 3.6+
• GPU: 推荐使用 NVIDIA GPU 以加速训练（可选，支持 CPU 运行）

前置依赖

请确保已安装以下核心库：

• torch >= 1.0.1
• torchvision >= 0.2.2
• tqdm == 4.31.1

国内加速建议：建议使用清华源或阿里源安装依赖，以提升下载速度。

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch torchvision tqdm

安装步骤

本项目主要为单文件实现，无需复杂的 pip install 包安装。请直接克隆仓库或将核心文件 attention_augmented_conv.py 下载到你的项目目录中。

1. 克隆仓库

   git clone https://github.com/leaderj1001/Attention-Augmented-Conv2d.git
   cd Attention-Augmented-Conv2d
   

2. 验证环境 确保 attention_augmented_conv.py 文件位于你的 Python 路径或当前工作目录下，以便直接导入。

基本使用

以下是使用 AugmentedConv 替换标准 nn.Conv2d 的最简示例。该模块支持与 CUDA 自动适配。

示例代码

import torch
from attention_augmented_conv import AugmentedConv

# 自动检测并使用 GPU
use_cuda = torch.cuda.is_available()
device = torch.device('cuda' if use_cuda else 'cpu')

# 创建随机输入张量 (Batch=16, Channels=3, Height=32, Width=32)
tmp = torch.randn((16, 3, 32, 32)).to(device)

# 初始化注意力增强卷积层
# 关键参数说明:
# in_channels/out_channels: 输入/输出通道数
# kernel_size: 卷积核大小
# dk: Key 的总维度 (必须是 Nh 的倍数)
# dv: Value 的总维度 (必须是 Nh 的倍数)
# Nh: 注意力头数量
# relative: 是否使用相对位置编码
# stride: 步长 (仅支持 1 或 2)
augmented_conv = AugmentedConv(
    in_channels=3, 
    out_channels=20, 
    kernel_size=3, 
    dk=40, 
    dv=4, 
    Nh=4, 
    relative=True, 
    stride=1, 
    shape=32  # 当 relative=True 时，需指定输入形状；若 stride>1，shape 应为输出形状
).to(device)

# 前向传播
conv_out = augmented_conv(tmp)

print(f"输出形状：{conv_out.shape}") 
# 预期输出：torch.Size([16, 20, 32, 32])

# 查看可学习参数（确认 key_rel_w 和 key_rel_h 是否存在）
for name, param in augmented_conv.named_parameters():
    print(f'parameter name: {name}')

关键注意事项

1. 参数整除约束：dk 和 dv 必须能被 Nh 整除，否则会触发断言错误。
2. Relative 模式下的 Shape：
   • 当 relative=True 时，必须提供 shape 参数。
   • 若 stride=1，shape 等于输入特征图尺寸。
   • 若 stride=2，shape 应等于输出特征图尺寸（即输入尺寸的一半）。
3. 步长限制：目前 stride 仅支持 1 或 2。