CBAM-TensorFlow
This is a Tensorflow implementation of “CBAM: Convolutional Block Attention Module”.
This repository includes the implementation of “Squeeze-and-Excitation Networks” as well, so that you can train and compare among base CNN model, base model with CBAM block and base model with SE block.
Base CNN models are ResNext, Inception-V4, and Inception-ResNet-V2 where the implementation is revised from Junho Kim’s code: SENet-Tensorflow.
If you want to use more sophisticated implementation and more base models to use, check the repository “CBAM-TensorFlow-Slim” which aims to be compatible on the TensorFlow-Slim image classification model library and support more base models.
CBAM: Convolutional Block Attention Module
CBAM proposes an architectural unit called “Convolutional Block Attention Module” (CBAM) block to improve representation power by using attention mechanism: focusing on important features and supressing unnecessary ones.
This research can be considered as a descendant and an improvement of “Squeeze-and-Excitation Networks”.
Diagram of a CBAM_block
Diagram of each attention sub-module
Classification results on ImageNet-1K
Prerequisites
- Python 3.x
- TensorFlow 1.x
- tflearn
Prepare Data set
This repository use Cifar10 dataset.
When you run the training script, the dataset will be automatically downloaded.
CBAM_block and SE_block Supportive Models
You can train and test base CNN model, base model with CBAM block and base model with SE block.
You can run CBAM_block or SE_block added models in the below list by adding one argument --attention_module=cbam_block or --attention_module=se_block when you train a model.
- Inception V4 + CBAM / + SE
- Inception-ResNet-v2 + CBAM / + SE
- ResNeXt + CBAM / + SE
Change Reduction ratio
To change reduction ratio, you can add an argument --reduction_ratio=8.
Train a Model
You can simply run a model by executing following scripts.
sh train_ResNext.shsh train_inception_resnet_v2.shsh train_inception_v4.sh
Train a model with CBAM_block
Below script gives you an example of training a model with CBAM_block.
CUDA_VISIBLE_DEVICES=0 python ResNeXt.py \
--model_name put_your_model_name \
--attention_module cbam_block \
--reduction_ratio 8 \
--learning_rate 0.1 \
--weight_decay 0.0005 \
--momentum 0.9 \
--batch_size 128 \
--total_epoch 100 \
--attention_module cbam_block
Train a model with SE_block
Below script gives you an example of training a model with SE_block.
CUDA_VISIBLE_DEVICES=0 python ResNeXt.py \
--model_name put_your_model_name \
--attention_module cbam_block \
--reduction_ratio 8 \
--learning_rate 0.1 \
--weight_decay 0.0005 \
--momentum 0.9 \
--batch_size 128 \
--total_epoch 100 \
--attention_module se_block
Train a model without attention module
Below script gives you an example of training a model without attention module.
CUDA_VISIBLE_DEVICES=0 python ResNeXt.py \
--model_name put_your_model_name \
--attention_module cbam_block \
--reduction_ratio 8 \
--learning_rate 0.1 \
--weight_decay 0.0005 \
--momentum 0.9 \
--batch_size 128 \
--total_epoch 100
Reference
Author
Byung Soo Ko / [email protected]
评论