CBAM.PyTorch

论文《CBAM：卷积块注意力模块》的非官方实现

简介

这些代码是论文《CBAM：卷积块注意力模块》的 PyTorch 重新实现版本。

Woo S, Park J, Lee J Y, et al. CBAM: Convolutional Block Attention Module[J]. 2018. ECCV2018

结构

CBAM 的整体结构。该模块包含两个顺序连接的子模块：通道注意力模块和空间注意力模块。在深度网络的每个卷积块之后，中间特征图都会通过我们的模块（CBAM）进行自适应地细化。

需求

• Python3
• PyTorch 0.4.1
• tensorboardX（可选）
• torchnet
• pretrainedmodels（可选）

结果

我们仅在 ImageNet-1K 数据集上测试了四种模型。训练集和验证集都缩放到 256 像素（短边），训练时仅使用 镜像翻转 和 随机裁剪 作为数据增强手段。验证时则采用中心裁剪获取 224x224 的图像块。

ImageNet-1K

模型	验证集（Top-1）	验证集（Top-5）
ResNet50	74.26	91.91
ResNet50-CBAM	75.45	92.55

CBAM.PyTorch 快速上手指南

CBAM (Convolutional Block Attention Module) 是一种轻量级的注意力模块，可串行集成到任何卷积神经网络中。该仓库提供了基于 PyTorch 的非官方复现版本，旨在通过通道注意力和空间注意力机制自适应地优化特征图，提升模型性能。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux / macOS / Windows
• Python 版本: Python 3.x
• 核心框架: PyTorch >= 0.4.1
• 可选依赖:
  • tensorboardX: 用于可视化训练过程
  • torchnet: 辅助工具库
  • pretrainedmodels: 用于加载预训练模型

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

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple <package_name>

安装步骤

1. 克隆项目代码 将仓库克隆到本地目录：

   git clone https://github.com/luuuyi/CBAM.PyTorch.git
   cd CBAM.PyTorch
   

2. 安装依赖包 使用 pip 安装必要的 Python 库（建议使用国内镜像）：

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

   注：如果您的环境中已安装 PyTorch，可跳过 torch 和 torchvision 的安装。

基本使用

CBAM 模块设计为即插即用，可以轻松嵌入到现有的 ResNet 或其他卷积网络结构中。以下是如何在代码中定义并集成 CBAM 的最简示例。

1. 导入模块

假设您已将 cbam.py（包含 CBAM 类定义的文件）放置在项目目录中：

from cbam import CBAM

2. 集成到模型中

您可以在任意卷积层之后添加 CBAM 模块。以下是一个在 ResNet 风格块中使用的示例片段：

import torch
import torch.nn as nn
from cbam import CBAM

class BasicBlockWithCBAM(nn.Module):
    def __init__(self, in_planes, out_planes, stride=1):
        super(BasicBlockWithCBAM, self).__init__()
        
        # 标准卷积层
        self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_planes)
        self.relu = nn.ReLU(inplace=True)
        
        self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_planes)

        # 集成 CBAM 模块
        # gate_channels: 通道数，reduction_ratio: 压缩比 (通常为 16)，pool_type: 池化类型
        self.cbam = CBAM(gate_channels=out_planes, reduction_ratio=16, pool_types=['avg', 'max'])

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        # 应用注意力机制
        out = self.cbam(out)

        out += residual
        out = self.relu(out)

        return out

3. 验证运行

创建一个简单的输入张量测试模块是否正常工作：

if __name__ == '__main__':
    # 模拟输入：Batch Size=4, Channels=64, Height=56, Width=56
    x = torch.randn(4, 64, 56, 56)
    
    # 初始化模块
    model = CBAM(gate_channels=64)
    
    # 前向传播
    output = model(x)
    
    print(f"Input shape: {x.shape}")
    print(f"Output shape: {output.shape}")

根据论文实验结果，在 ImageNet-1K 数据集上，ResNet50 集成 CBAM 后 Top-1 准确率可从 74.26% 提升至 75.45%。