BAM 和 CBAM

官方 PyTorch 代码，对应于“BAM：瓶颈注意力模块（BMVC2018）”和“CBAM：卷积块注意力模块（ECCV2018）”。

更新与通知

• 2018-10-08：目前仅验证了 CBAM 的测试代码。训练代码可能存在细微错误。将在几天内修复。
• 2018-10-11：训练代码已验证完毕。新增了 RESNET50+BAM 预训练权重。

环境要求

代码已在以下环境中验证通过：

• Ubuntu 16.04，4 块 GTX 1080 Ti 显卡，Docker（PyTorch 0.4.1，CUDA 9.0 + CuDNN 7.0，Python 3.6）

使用方法

包含基于 ResNet50 的示例。示例脚本位于 ./scripts/ 目录下。 ImageNet 数据应放置在 ./data/ImageNet/ 目录下，并分别创建名为 train 和 val 的子文件夹。

# 使用 BAM 训练（ResNet50 为骨干网络）
python train_imagenet.py --ngpu 4 --workers 20 --arch resnet --depth 50 --epochs 100 --batch-size 256 --lr 0.1 --att-type BAM --prefix RESNET50_IMAGENET_BAM ./data/ImageNet
# 使用 CBAM 训练（ResNet50 为骨干网络）
python train_imagenet.py --ngpu 4 --workers 20 --arch resnet --depth 50 --epochs 100 --batch-size 256 --lr 0.1 --att-type CBAM --prefix RESNET50_IMAGENET_CBAM ./data/ImageNet

从检查点继续训练

• ResNet50+CBAM（训练 100 个 epoch）的检查点已提供在此 链接 中。ACC@1=77.622，ACC@5=93.948。
• ResNet50+BAM（训练 90 个 epoch）的检查点已提供在此 链接 中。ACC@1=76.860，ACC@5=93.416。

验证时，请使用如下脚本：

python train_imagenet.py --ngpu 4 --workers 20 --arch resnet --depth 50 --att-type CBAM --prefix EVAL --resume $CHECKPOINT_PATH$ --evaluate ./data/ImageNet

其他实现

• bruinxiong 提供的 带有若干修改的 MXNet 版本 CBAM 实现

attention-module 快速上手指南

本指南基于官方 README 整理，帮助开发者快速在 PyTorch 环境中集成 BAM (Bottleneck Attention Module) 和 CBAM (Convolutional Block Attention Module) 注意力模块。

环境准备

代码已在以下环境中验证通过，建议参考此配置搭建开发环境：

• 操作系统: Ubuntu 16.04
• 硬件要求: 推荐多卡环境（原文使用 4*GTX 1080 Ti）
• 容器化: 推荐使用 Docker
• 核心依赖:
  • PyTorch 0.4.1
  • CUDA 9.0 + CuDNN 7.0
  • Python 3.6

注意：虽然原文基于较旧版本的 PyTorch，但注意力模块的核心代码通常具有较好的兼容性。若使用新版 PyTorch，可能需微调部分 API。

安装步骤

本项目为纯代码实现，无需通过 pip 安装额外包，只需克隆仓库并准备数据即可。

1. 克隆项目代码

   git clone https://github.com/Jongchan/attention-module.git
   cd attention-module
   

2. 准备数据集 将 ImageNet 数据集放置在 ./data/ImageNet/ 目录下，并确保包含以下子文件夹结构：

   ./data/ImageNet/
   ├── train/
   └── val/
   

基本使用

项目提供了基于 ResNet50 的训练脚本。你可以通过修改 --att-type 参数来选择使用 BAM 或 CBAM 模块。

1. 开始训练

使用 BAM 模块训练 (ResNet50 backbone):

python train_imagenet.py --ngpu 4 --workers 20 --arch resnet --depth 50 --epochs 100 --batch-size 256 --lr 0.1 --att-type BAM --prefix RESNET50_IMAGENET_BAM ./data/ImageNet

使用 CBAM 模块训练 (ResNet50 backbone):

python train_imagenet.py --ngpu 4 --workers 20 --arch resnet --depth 50 --epochs 100 --batch-size 256 --lr 0.1 --att-type CBAM --prefix RESNET50_IMAGENET_CBAM ./data/ImageNet

2. 加载预训练模型进行评估

官方提供了训练好的检查点（Checkpoints），你可以下载后直接进行验证评估。

• ResNet50+CBAM (100 epochs): Top-1 Acc 77.62%
• ResNet50+BAM (90 epochs): Top-1 Acc 76.86%

下载对应的 checkpoint 文件后，运行以下命令进行验证（请将 $CHECKPOINT_PATH$ 替换为实际文件路径）：

python train_imagenet.py --ngpu 4 --workers 20 --arch resnet --depth 50 --att-type CBAM --prefix EVAL --resume $CHECKPOINT_PATH$ --evaluate ./data/ImageNet