别再死记硬背了！用PyTorch代码实战理解5大2D注意力机制（附Non-Local/SE/CBAM对比）

用PyTorch代码实战理解5大2D注意力机制：从Non-Local到Dual-Attention

在深度学习领域，注意力机制已经成为提升模型性能的关键技术。但对于初学者来说，理论公式往往让人望而生畏。本文将带你用PyTorch代码实现5种主流2D注意力机制，通过可视化特征图和修改参数来直观理解它们的工作原理。我们将在CIFAR-10数据集上对比Non-Local、SE、CBAM等模块的效果，让你真正掌握如何将这些技术应用到自己的项目中。

1. 准备工作与环境搭建

在开始实现注意力机制前，我们需要搭建一个基础实验环境。这里使用PyTorch 1.10+和Torchvision，建议在Python 3.8+环境中运行以下代码：

import torch import torch.nn as nn import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader import matplotlib.pyplot as plt # 检查GPU可用性 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # 数据预处理 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) # 加载CIFAR-10数据集 train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) train_loader = DataLoader(train_set, batch_size=64, shuffle=True) test_loader = DataLoader(test_set, batch_size=64, shuffle=False)

为了评估不同注意力机制的效果，我们定义一个基础ResNet模型作为backbone：

class BasicBlock(nn.Module): expansion = 1 def __init__(self, in_channels, out_channels, stride=1): super().__init__() self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(out_channels) self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(out_channels) self.shortcut = nn.Sequential() if stride != 1 or in_channels != self.expansion*out_channels: self.shortcut = nn.Sequential( nn.Conv2d(in_channels, self.expansion*out_channels, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(self.expansion*out_channels) ) def forward(self, x): out = torch.relu(self.bn1(self.conv1(x))) out = self.bn2(self.conv2(out)) out += self.shortcut(x) out = torch.relu(out) return out

2. Non-Local注意力机制实现与解析

Non-Local是一种捕捉长距离依赖关系的注意力机制，特别适合处理需要全局上下文信息的任务。让我们先看PyTorch实现：

class NonLocalBlock(nn.Module): def __init__(self, in_channels, inter_channels=None): super().__init__() self.in_channels = in_channels self.inter_channels = inter_channels if inter_channels else in_channels // 2 self.g = nn.Conv2d(in_channels, self.inter_channels, kernel_size=1) self.theta = nn.Conv2d(in_channels, self.inter_channels, kernel_size=1) self.phi = nn.Conv2d(in_channels, self.inter_channels, kernel_size=1) self.W = nn.Conv2d(self.inter_channels, in_channels, kernel_size=1) self.W.weight.data.zero_() self.W.bias.data.zero_() def forward(self, x): batch_size = x.size(0) g_x = self.g(x).view(batch_size, self.inter_channels, -1) g_x = g_x.permute(0, 2, 1) theta_x = self.theta(x).view(batch_size, self.inter_channels, -1) theta_x = theta_x.permute(0, 2, 1) phi_x = self.phi(x).view(batch_size, self.inter_channels, -1) f = torch.matmul(theta_x, phi_x) f = torch.softmax(f, dim=-1) y = torch.matmul(f, g_x) y = y.permute(0, 2, 1).contiguous() y = y.view(batch_size, self.inter_channels, *x.size()[2:]) y = self.W(y) z = y + x return z

关键点解析：

• g,theta,phi三个1x1卷积分别生成query、key和value
• 通过矩阵乘法计算注意力权重，softmax归一化
• 最终输出是原始输入与注意力加权的特征相加

我们可以将这个模块插入到ResNet中：

class ResNetWithNonLocal(nn.Module): def __init__(self, block, num_blocks, num_classes=10): super().__init__() self.in_channels = 64 self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(64) self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1) self.nonlocal1 = NonLocalBlock(64) self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2) self.nonlocal2 = NonLocalBlock(128) self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2) self.nonlocal3 = NonLocalBlock(256) self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2) self.linear = nn.Linear(512*block.expansion, num_classes) def _make_layer(self, block, out_channels, num_blocks, stride): # ... (与标准ResNet实现相同) def forward(self, x): out = torch.relu(self.bn1(self.conv1(x))) out = self.layer1(out) out = self.nonlocal1(out) out = self.layer2(out) out = self.nonlocal2(out) out = self.layer3(out) out = self.nonlocal3(out) out = self.layer4(out) out = torch.avg_pool2d(out, 4) out = out.view(out.size(0), -1) out = self.linear(out) return out

提示：Non-Local模块计算开销较大，实际使用时可以考虑在高层特征图（分辨率较低）上应用，或者在theta和phi后添加下采样操作。

3. SE（Squeeze-and-Excitation）模块实现

SE模块通过显式建模通道间关系来自适应地重新校准通道特征响应。下面是其PyTorch实现：

class SEBlock(nn.Module): def __init__(self, in_channels, reduction=16): super().__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Sequential( nn.Linear(in_channels, in_channels // reduction), nn.ReLU(inplace=True), nn.Linear(in_channels // reduction, in_channels), nn.Sigmoid() ) def forward(self, x): b, c, _, _ = x.size() y = self.avg_pool(x).view(b, c) y = self.fc(y).view(b, c, 1, 1) return x * y.expand_as(x)

SE模块特点：

• 先通过全局平均池化压缩空间信息（Squeeze）
• 然后通过两个全连接层学习通道间依赖关系（Excitation）
• 最后将学习到的权重应用到原始特征上

将SE模块整合到ResNet中的示例：

class SEBasicBlock(nn.Module): expansion = 1 def __init__(self, in_channels, out_channels, stride=1, reduction=16): super().__init__() self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(out_channels) self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(out_channels) self.se = SEBlock(out_channels, reduction) self.shortcut = nn.Sequential() if stride != 1 or in_channels != self.expansion*out_channels: self.shortcut = nn.Sequential( nn.Conv2d(in_channels, self.expansion*out_channels, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(self.expansion*out_channels) ) def forward(self, x): out = torch.relu(self.bn1(self.conv1(x))) out = self.bn2(self.conv2(out)) out = self.se(out) out += self.shortcut(x) out = torch.relu(out) return out

4. CBAM（Convolutional Block Attention Module）实现

CBAM结合了通道注意力和空间注意力，是一种轻量级但有效的注意力模块。下面是完整实现：

class ChannelAttention(nn.Module): def __init__(self, in_channels, reduction=16): super().__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) self.fc = nn.Sequential( nn.Conv2d(in_channels, in_channels // reduction, 1, bias=False), nn.ReLU(), nn.Conv2d(in_channels // reduction, in_channels, 1, bias=False) ) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = self.fc(self.avg_pool(x)) max_out = self.fc(self.max_pool(x)) out = avg_out + max_out return self.sigmoid(out) class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super().__init__() self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) x = torch.cat([avg_out, max_out], dim=1) x = self.conv(x) return self.sigmoid(x) class CBAM(nn.Module): def __init__(self, in_channels, reduction=16, kernel_size=7): super().__init__() self.ca = ChannelAttention(in_channels, reduction) self.sa = SpatialAttention(kernel_size) def forward(self, x): x = x * self.ca(x) x = x * self.sa(x) return x

CBAM模块特点：

• 通道注意力分支同时考虑平均池化和最大池化信息
• 空间注意力分支通过卷积操作学习空间位置的重要性
• 两个注意力分支顺序应用，先通道后空间

将CBAM整合到ResNet中的示例：

class CBAMBasicBlock(nn.Module): expansion = 1 def __init__(self, in_channels, out_channels, stride=1, reduction=16): super().__init__() self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(out_channels) self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(out_channels) self.cbam = CBAM(out_channels, reduction) self.shortcut = nn.Sequential() if stride != 1 or in_channels != self.expansion*out_channels: self.shortcut = nn.Sequential( nn.Conv2d(in_channels, self.expansion*out_channels, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(self.expansion*out_channels) ) def forward(self, x): out = torch.relu(self.bn1(self.conv1(x))) out = self.bn2(self.conv2(out)) out = self.cbam(out) out += self.shortcut(x) out = torch.relu(out) return out

5. Dual-Attention与Criss-Cross注意力实现

5.1 Dual-Attention网络

Dual-Attention同时考虑位置注意力和通道注意力，特别适合语义分割等密集预测任务：

class PositionAttention(nn.Module): def __init__(self, in_channels): super().__init__() self.conv_q = nn.Conv2d(in_channels, in_channels // 8, 1) self.conv_k = nn.Conv2d(in_channels, in_channels // 8, 1) self.conv_v = nn.Conv2d(in_channels, in_channels, 1) self.gamma = nn.Parameter(torch.zeros(1)) self.softmax = nn.Softmax(dim=-1) def forward(self, x): b, c, h, w = x.size() q = self.conv_q(x).view(b, -1, h*w).permute(0, 2, 1) k = self.conv_k(x).view(b, -1, h*w) v = self.conv_v(x).view(b, -1, h*w) attn = torch.bmm(q, k) attn = self.softmax(attn) out = torch.bmm(v, attn.permute(0, 2, 1)) out = out.view(b, c, h, w) return self.gamma * out + x class ChannelAttention(nn.Module): def __init__(self, in_channels): super().__init__() self.gamma = nn.Parameter(torch.zeros(1)) self.softmax = nn.Softmax(dim=-1) def forward(self, x): b, c, h, w = x.size() q = x.view(b, c, -1) k = x.view(b, c, -1).permute(0, 2, 1) v = x.view(b, c, -1) attn = torch.bmm(q, k) attn = self.softmax(attn) out = torch.bmm(attn, v) out = out.view(b, c, h, w) return self.gamma * out + x class DualAttention(nn.Module): def __init__(self, in_channels): super().__init__() self.position = PositionAttention(in_channels) self.channel = ChannelAttention(in_channels) def forward(self, x): p_out = self.position(x) c_out = self.channel(x) return p_out + c_out

5.2 Criss-Cross注意力

Criss-Cross注意力通过交叉路径捕获上下文信息，计算效率比Non-Local更高：

class CrissCrossAttention(nn.Module): def __init__(self, in_channels): super().__init__() self.conv_q = nn.Conv2d(in_channels, in_channels // 8, 1) self.conv_k = nn.Conv2d(in_channels, in_channels // 8, 1) self.conv_v = nn.Conv2d(in_channels, in_channels, 1) self.gamma = nn.Parameter(torch.zeros(1)) self.softmax = nn.Softmax(dim=3) def forward(self, x): b, c, h, w = x.size() q = self.conv_q(x) # [b, c', h, w] k = self.conv_k(x) # [b, c', h, w] v = self.conv_v(x) # [b, c, h, w] # 水平方向注意力 q_h = q.permute(0, 3, 1, 2).contiguous().view(b*w, -1, h) k_h = k.permute(0, 3, 1, 2).contiguous().view(b*w, -1, h) v_h = v.permute(0, 3, 1, 2).contiguous().view(b*w, -1, h) attn_h = torch.bmm(q_h.permute(0, 2, 1), k_h) attn_h = self.softmax(attn_h) out_h = torch.bmm(v_h, attn_h.permute(0, 2, 1)) out_h = out_h.view(b, w, -1, h).permute(0, 2, 3, 1) # 垂直方向注意力 q_v = q.permute(0, 2, 1, 3).contiguous().view(b*h, -1, w) k_v = k.permute(0, 2, 1, 3).contiguous().view(b*h, -1, w) v_v = v.permute(0, 2, 1, 3).contiguous().view(b*h, -1, w) attn_v = torch.bmm(q_v.permute(0, 2, 1), k_v) attn_v = self.softmax(attn_v) out_v = torch.bmm(v_v, attn_v.permute(0, 2, 1)) out_v = out_v.view(b, h, -1, w).permute(0, 2, 1, 3) out = self.gamma * (out_h + out_v) + x return out

6. 注意力机制对比与实验分析

为了比较不同注意力机制的效果，我们在CIFAR-10上进行了实验。以下是实验结果对比：

各注意力机制特点总结：

1. SE模块：

   • 计算量小，易于集成
   • 主要关注通道间关系
   • 适合资源受限场景

2. CBAM：

   • 同时考虑通道和空间注意力
   • 计算开销适中
   • 通用性强，适合大多数视觉任务

3. Non-Local：

   • 捕获长距离依赖关系
   • 计算开销大
   • 适合需要全局上下文的场景

4. Dual-Attention：

   • 位置和通道注意力并行
   • 性能提升明显但计算量较大
   • 适合密集预测任务

5. Criss-Cross：

   • 交叉路径捕获上下文
   • 比Non-Local更高效
   • 特别适合语义分割

注意：在实际项目中，选择注意力机制时需要权衡计算开销和性能提升。对于计算资源有限的场景，SE或CBAM通常是更好的选择；而对于需要捕获长距离依赖的任务，Non-Local或Dual-Attention可能更合适。

7. 注意力机制可视化与调试技巧

理解注意力机制最直观的方式是可视化其激活图。以下代码展示了如何可视化CBAM模块的注意力权重：

def visualize_attention(model, image): # 前向传播并获取中间层输出 activations = {} def hook_fn(module, input, output): activations[module._get_name()] = output.detach() hooks = [] for name, module in model.named_modules(): if isinstance(module, (ChannelAttention, SpatialAttention)): hooks.append(module.register_forward_hook(hook_fn)) with torch.no_grad(): model(image.unsqueeze(0).to(device)) # 移除hooks for hook in hooks: hook.remove() # 可视化 fig, axes = plt.subplots(1, 3, figsize=(15, 5)) axes[0].imshow(image.permute(1, 2, 0).cpu().numpy()) axes[0].set_title('Original Image') axes[0].axis('off') if 'ChannelAttention' in activations: channel_attn = activations['ChannelAttention'].squeeze().cpu().numpy() axes[1].barh(range(len(channel_attn)), channel_attn) axes[1].set_title('Channel Attention Weights') if 'SpatialAttention' in activations: spatial_attn = activations['SpatialAttention'].squeeze().cpu().numpy() axes[2].imshow(spatial_attn, cmap='hot') axes[2].set_title('Spatial Attention Map') axes[2].axis('off') plt.tight_layout() plt.show()

调试注意力网络的实用技巧：

1. 初始化策略：

   • 注意力权重初始化为接近零的小值，让网络初期主要依赖原始特征
   • 例如：self.gamma = nn.Parameter(torch.zeros(1))

2. 计算优化：

   • 对大分辨率特征图，先在theta和phi路径添加下采样
   • 使用分组卷积或深度可分离卷积减少计算量

3. 训练技巧：

   • 初始阶段可以固定backbone，只训练注意力模块
   • 使用渐进式训练策略，逐步引入注意力模块

4. 常见问题排查：

   • 如果性能没有提升，检查注意力权重是否过于均匀（没有学到有用模式）
   • 监控注意力权重的分布，避免出现极端值（全0或全1）

# 监控注意力权重分布的代码示例 def monitor_attention_distribution(model, dataloader): model.eval() attn_weights = [] with torch.no_grad(): for images, _ in dataloader: outputs = model(images.to(device)) # 假设模型会返回注意力权重 if hasattr(model, 'get_attention_weights'): weights = model.get_attention_weights() attn_weights.append(weights.cpu()) attn_weights = torch.cat(attn_weights) plt.hist(attn_weights.numpy().flatten(), bins=50) plt.xlabel('Attention Weight Value') plt.ylabel('Frequency') plt.title('Attention Weights Distribution') plt.show()