CBAM神经网络实战：如何解决注意力机制中的特征选择难题

1次阅读

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

在计算机视觉任务中，传统的卷积神经网络 (CNN) 存在一个明显的局限性：它们对所有区域的特征都给予同等重视。全局平均池化 (Global Average Pooling) 虽然简单高效，但会导致空间信息的丢失。这就好比看一幅画时，我们人类会自然地聚焦在关键区域，而 CNN 却试图平等地关注每一个像素。

目前主流的注意力机制包括：

SE 模块(Squeeze-and-Excitation)：仅关注通道维度，计算量小但忽略空间关系
Non-local 网络：捕获长距离依赖，但计算复杂度高(O(N^2))
CBAM(Convolutional Block Attention Module)：结合通道和空间注意力，平衡效果和效率

量化指标对比(V100 32GB, ImageNet-1k):

方法	参数量(M)	GFLOPs	Top-1 Acc
ResNet50	25.5	4.1	76.2%
+SE	26.3	4.2	77.1%
+Non-local	28.7	5.9	77.5%
+CBAM	26.8	4.3	78.5%

import torch
import torch.nn as nn

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        # 1x1 卷积降维，减少计算量
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.fc1 = nn.Conv2d(in_planes, in_planes//ratio, 1, bias=False)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Conv2d(in_planes//ratio, in_planes, 1, bias=False)

        self.sigmoid = nn.Sigmoid()  # 将权重归一化到 0 -1

    def forward(self, x):
        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmoid(out)

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()

        assert kernel_size in (3,7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, 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.conv1(x)
        return self.sigmoid(x)

class CBAM(nn.Module):
    def __init__(self, in_planes, ratio=16, kernel_size=7):
        super(CBAM, self).__init__()
        self.ca = ChannelAttention(in_planes, ratio)
        self.sa = SpatialAttention(kernel_size)

    def forward(self, x):
        x = x * self.ca(x)  # 通道注意力
        x = x * self.sa(x)  # 空间注意力
        return x

Nsight 分析：使用 NVIDIA Nsight 工具分析发现，注意力模块的矩阵乘是主要计算热点。解决方案：
将小矩阵乘法替换为分组卷积
使用混合精度训练(AMP)
TensorRT 部署：
将连续的 1 ×1 卷积和 BN 层融合
使用 trt.NetworkDefinition 的add_plugin_v2接口自定义注意力算子

BN 与注意力协同：在残差分支中添加 CBAM 时，建议将 BN 层放在注意力模块之前，避免破坏统计特性
多 GPU 训练 ：当 batch size 较小时，使用torch.nn.parallel.DistributedDataParallel 替代DataParallel，确保梯度同步准确

CBAM 的注意力机制在 2D 图像上表现出色，但如何将其扩展到 3D 点云处理？可能的思路：
1. 将通道注意力扩展到特征维度
2. 用球状卷积替代空间注意力中的 2D 卷积
3. 考虑点云的不规则特性，设计自适应的局部注意力窗口

在实际工业场景中，我们还需要权衡计算成本和精度提升。CBAM 虽然增加了少量计算量，但通过合理的优化和部署，完全可以在嵌入式设备上实时运行。

正文完