This is a Tensorflow implementation of "Squeeze-and-Excitation Networks" aiming to be compatible on the TensorFlow-Slim image classification model library.
SENet proposes an architectural unit called "Squeeze-and-Excitation" (SE) block to improve the representational power of a network by explicitly modelling the interdependencies between the channels of its convolutional features.
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 SE-block added models in the below list by adding one argument --attention_module='se_block'
when you train or evaluate a model.
To change reduction ratio, you have to manually set the ratio on def se_block(residual, name, ratio=8)
method in SENet-tensorflow-slim/nets/attention_module.py
.
Below script gives you an example of training a model with SE-block.
Don't forget to put an argument --attention_module=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 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
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 SE-block during training.
Don't forget to put an argument --attention_module=se_block
.
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 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
Byung Soo Ko / kobiso62@gmail.com