This is a Tensorflow implementation of "CBAM: Convolutional Block Attention Module" aiming to be compatible on the TensorFlow-Slim image classification model library. This repository includes the implementation of "SENet-tensorflow-slim".
If you want to use simpler implementation, check the repository "CBAM-TensorFlow" which includes simple tensorflow implementation of ResNext, Inception-V4, and Inception-ResNet-V2 on Cifar10 dataset.
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".
- Python 3.x
- TensorFlow 1.x
- TF-slim
You should prepare your own dataset or open dataset (Cifar10, flowers, MNIST, ImageNet). For preparing dataset, you can follow the 'preparing the datasets' part in TF-Slim image models README.
This project is based on TensorFlow-Slim image classification model library.
Every image classification model in TensorFlow-Slim can be run the same.
And, 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 or evaluate a model.
- Inception V4 + CBAM / + SE
- Inception-ResNet-v2 + CBAM / + SE
- ResNet V1 50 + CBAM / + SE
- ResNet V1 101 + CBAM / + SE
- ResNet V1 152 + CBAM / + SE
- ResNet V1 200 + CBAM / + SE
- ResNet V2 50 + CBAM / + SE
- ResNet V2 101 + CBAM / + SE
- ResNet V2 152 + CBAM / + SE
- ResNet V2 200 + CBAM / + SE
To change reduction ratio, you have to manually set the ratio on def cbam_block(input_feature, name, ratio=16)
method for cbam_block or def se_block(residual, name, ratio=8)
method for se_block in CBAM-tensorflow-slim/nets/attention_module.py
.
You can find example of training script in CBAM-tensorflow-slim/scripts/
.
Below script gives you an example of training a model with CBAM_block.
DATASET_DIR=/DIRECTORY/TO/DATASET
TRAIN_DIR=/DIRECTORY/TO/TRAIN
CUDA_VISIBLE_DEVICES=0 python train_image_classifier.py \
--train_dir=${TRAIN_DIR} \
--dataset_name=imagenet \
--dataset_split_name=train \
--dataset_dir=${DATASET_DIR} \
--model_name=resnet_v1_50 \
--batch_size=100 \
--attention_module=cbam_block
Below script gives you an example of training a model with SE_block.
DATASET_DIR=/DIRECTORY/TO/DATASET
TRAIN_DIR=/DIRECTORY/TO/TRAIN
CUDA_VISIBLE_DEVICES=0 python train_image_classifier.py \
--train_dir=${TRAIN_DIR} \
--dataset_name=imagenet \
--dataset_split_name=train \
--dataset_dir=${DATASET_DIR} \
--model_name=resnet_v1_50 \
--batch_size=100 \
--attention_module=se_block
Below script gives you an example of training a model without attention module.
DATASET_DIR=/DIRECTORY/TO/DATASET
TRAIN_DIR=/DIRECTORY/TO/TRAIN
CUDA_VISIBLE_DEVICES=0 python train_image_classifier.py \
--train_dir=${TRAIN_DIR} \
--dataset_name=imagenet \
--dataset_split_name=train \
--dataset_dir=${DATASET_DIR} \
--model_name=resnet_v1_50 \
--batch_size=100
You can find example of evaluation script in CBAM-tensorflow-slim/scripts/
.
To keep track of validation accuracy while training, you can use eval_image_classifier_loop.py
which evaluate the performance at multiple checkpoints during training.
If you want to just evaluate a model once, you can use eval_image_classifier.py
.
Below script gives you an example of evaluating a model with CBAM_block during training.
DATASET_DIR=/DIRECTORY/TO/DATASET
CHECKPOINT_FILE=/DIRECTORY/TO/CHECKPOINT
EVAL_DIR=/DIRECTORY/TO/EVAL
CUDA_VISIBLE_DEVICES=0 python eval_image_classifier_loop.py \
--alsologtostderr \
--checkpoint_path=${CHECKPOINT_FILE} \
--dataset_dir=${DATASET_DIR} \
--eval_dir=${EVAL_DIR} \
--dataset_name=imagenet \
--dataset_split_name=validation \
--model_name=resnet_v1_50 \
--batch_size=100 \
--attention_module=cbam_block
Below script gives you an example of evaluating a model with SE_block during training.
DATASET_DIR=/DIRECTORY/TO/DATASET
CHECKPOINT_FILE=/DIRECTORY/TO/CHECKPOINT
EVAL_DIR=/DIRECTORY/TO/EVAL
CUDA_VISIBLE_DEVICES=0 python eval_image_classifier_loop.py \
--alsologtostderr \
--checkpoint_path=${CHECKPOINT_FILE} \
--dataset_dir=${DATASET_DIR} \
--eval_dir=${EVAL_DIR} \
--dataset_name=imagenet \
--dataset_split_name=validation \
--model_name=resnet_v1_50 \
--batch_size=100 \
--attention_module=se_block
Below script gives you an example of evaluating a model without attention module during training.
DATASET_DIR=/DIRECTORY/TO/DATASET
CHECKPOINT_FILE=/DIRECTORY/TO/CHECKPOINT
EVAL_DIR=/DIRECTORY/TO/EVAL
CUDA_VISIBLE_DEVICES=0 python eval_image_classifier_loop.py \
--alsologtostderr \
--checkpoint_path=${CHECKPOINT_FILE} \
--dataset_dir=${DATASET_DIR} \
--eval_dir=${EVAL_DIR} \
--dataset_name=imagenet \
--dataset_split_name=validation \
--model_name=resnet_v1_50 \
--batch_size=100
- Blog: CBAM: Convolutional Block Attention Module
- Repository: CBAM-TensorFlow
- Repository: CBAM-Keras
- Repository: SENet-TensorFlow-Slim
- Paper: CBAM: Convolutional Block Attention Module
- Paper: Squeeze-and-Excitation Networks
- Repository: TensorFlow-Slim image classification model library
Byung Soo Ko / [email protected]