search

megvii-model / MABN

MIT License
182 stars 27 forks source link

readme

Moving Average Batch Normalization

This reposity is the Pytorch implementation of Moving Average Batch Normalization on Imagenet classfication, COCO object detection and instance segmentation tasks.

The paper has been published as an ICLR2020 conference paper (https://openreview.net/forum?id=SkgGjRVKDS&noteId=BJeCWt3KiH).

Results

Overall comparation of MABN and its counterparts

Top 1 Error versus Batch Size:

Inference Speend

Norm Iterations/second
BN/MABN 237.88
Instance Normalization 105.60
Group Normalization 99.37
Layer Normalization 125.44

Imagenet

Model Normalization Batch size Norm Top 1 Accuracy
ResNet50 32 BN 76.59
ResNet50 2 BN 64.78
ResNet50 2 BRN 69.71
ResNet50 2 MABN 76.33

COCO

Backbone Method Training Strategy Norm Batch Size APb APb0.50 APb0.75 APm APm0.50 APm0.75
R50-FPN Mask R-CNN 2x from scratch BN 2 32.38 50.44 35.47 29.07 47.68 30.75
R50-FPN Mask R-CNN 2x from scratch BRN 2 34.07 52.66 37.12 30.98 50.03 32.93
R50-FPN Mask R-CNN 2x from scratch SyncBN 2x8 36.80 56.06 40.23 33.10 53.15 35.24
R50-FPN Mask R-CNN 2x from scratch MABN 2 36.50 55.79 40.17 32.69 52.78 34.71
R50-FPN Mask R-CNN 2x fine-tune SyncBN 2x8 38.25 57.81 42.01 34.22 54.97 36.34
R50-FPN Mask R-CNN 2x fine-tune MABN 2 38.42 58.19 41.99 34.12 55.10 36.12

Usage

The formal implementation of MABN is in MABN.py. You can use MABN by directly replacing torch.nn.BatchNorm2d and torch.nn.Conv2d with MABN2d and CenConv2d respectively in your project. Please don't forget to set extra args in MABN2d.

Demo

One node with 8 GPUs.

Imagenet

Notice the Imagenet classification simulate the small batch training settings by using small normalization batch size and regular SGD batch size.

cd /your_path_to_repo/cls

# 8 GPUs Train and Test
python3 -m torch.distributed.launch --nproc_per_node=8 train.py --gpu_num=8 \
    --save /your_path_to_logs \
    --train_dir /your_imagenet_training_dataset_dir \
    --val_dir /your_imagenet_eval_dataset_dir \
    --gpu_num=8

# Only Test the trained model
python3 -m torch.distributed.launch --nproc_per_node=1 train.py --gpu_num=1 \
    --save /your_path_to_logs \
    --val_dir /your_imagenet_eval_dataset_dir \
    --checkpoint_dir /your_path_to_checkpoint \
    --test_only

COCO

Install

Please refer to INSTALL.md for installation and dataset preparation.

To use SyncBN, please do:

cd /your_path_to_repo/det/maskrcnn_benchmar/distributed_syncbn
bash compile.sh

Imagenet Pretrained model

You can download the pretrained model of ResNet-50 in:

Dropbox: https://www.dropbox.com/sh/fbsi6935vmatbi9/AAA2jv0EBcSgySTgZnNZ3lmPa?dl=0 ;

Baiduyun: https://pan.baidu.com/s/1Md_UzwWEiZZKu84R0yZ6aw password: zww2 .

Notice R-50-2.pkl is the for SyncBN, while R50-wc.pth is for MABN.

Run Experiment

cd /your_path_to_repo/det
# Train MABN from scratch
python3 -m torch.distributed.launch --nproc_per_node=8 tools/train_net.py \
    --skip-test \
    --config-file configs/e2e_mask_rcnn_R_50_FPN_mabn_2x_from_scratch.yaml \
    DATALOADER.NUM_WORKERS 2 \
    OUTPUT_DIR /your_path_to_logs

# Train MABN fine tuning (Download the pertrained model and set the path in configs at first)
python3 -m torch.distributed.launch --nproc_per_node=8 tools/train_net.py \
    --skip-test \
    --config-file configs/e2e_mask_rcnn_R_50_FPN_mabn_2x_fine_tune.yaml \
    DATALOADER.NUM_WORKERS 2 \
    OUTPUT_DIR /your_path_to_logs

# Test model
python3 -m torch.distributed.launch --nproc_per_node=8 tools/test_net.py \
       --config-file configs/e2e_mask_rcnn_R_50_FPN_mabn_2x_from_scratch.yaml \
       MODEL.WEIGHT /your_path_to_logs/model_0180000.pth \
       TEST.IMS_PER_BATCH 8

Thanks

This implementation of COCO is based on maskrcnn-benchmark. Ref to this link for more details about maskrcnn-benchmark.

Citation

If you use Moving Average Batch Normalization in your research, please cite:

@inproceedings{
yan2020towards,
title={Towards Stablizing Batch Statistics in Backward Propagation of Batch Normalization},
author={Junjie Yan, Ruosi Wan, Xiangyu Zhang, Wei Zhang, Yichen Wei, Jian Sun},
booktitle={International Conference on Learning Representations},
year={2020},
url={https://openreview.net/forum?id=SkgGjRVKDS}
}