EMAN: Exponential Moving Average Normalization for Self-supervised and Semi-supervised Learning

This is a PyTorch implementation of the EMAN paper. It supports three popular self-supervised and semi-supervised learning techniques, i.e., MoCo, BYOL and FixMatch.

If you use the code/model/results of this repository please cite:

@inproceedings{cai21eman,
  author  = {Zhaowei Cai and Avinash Ravichandran and Subhransu Maji and Charless Fowlkes and Zhuowen Tu and Stefano Soatto},
  title   = {Exponential Moving Average Normalization for Self-supervised and Semi-supervised Learning},
  booktitle = {CVPR},
  Year  = {2021}
}

Install

First, install PyTorch and torchvision. We have tested on version of 1.7.1, but the other versions should also be working, e.g. 1.5.1.

$ conda install pytorch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 cudatoolkit=10.1 -c pytorch

Also install other dependencies.

$ pip install pandas faiss-gpu

Data Preparation

Install ImageNet dataset following the official PyTorch ImageNet training code, with the standard data folder structure for the torchvision datasets.ImageFolder. Please download the ImageNet index files for semi-supervised learning experiments. The file structure should look like:

$ tree data
imagenet
├── train
│   ├── class1
│   │   └── *.jpeg
│   ├── class2
│   │   └── *.jpeg
│   └── ...
├── val
│   ├── class1
│   │   └── *.jpeg
│   ├── class2
│   │   └── *.jpeg
│   └── ...
└── indexes
    └── *_index.csv

Training

To do self-supervised pre-training of MoCo-v2 with EMAN for 200 epochs, run:

python main_moco.py \
  --arch MoCoEMAN --backbone resnet50_encoder \
  --epochs 200 --warmup-epoch 10 \
  --moco-t 0.2 --cos \
  --dist-url 'tcp://localhost:10001' --multiprocessing-distributed --world-size 1 --rank 0 \
  /path/to/imagenet

To do self-supervised pre-training of BYOL with EMAN for 200 epochs, run:

python main_byol.py \
  --arch BYOLEMAN --backbone resnet50_encoder \
  --lr 1.8 -b 512 --wd 0.000001 \
  --byol-m 0.98 \
  --epochs 200 --cos --warmup-epoch 10 \
  --dist-url 'tcp://localhost:10001' --multiprocessing-distributed --world-size 1 --rank 0 \
  /path/to/imagenet

To do semi-supervised training of FixMatch with EMAN for 100 epochs, run:

python main_fixmatch.py \
  --arch FixMatch --backbone resnet50_encoder \
  --eman \
  --lr 0.03 \
  --epochs 100 --schedule 60 80 \
  --warmup-epoch 5 \
  --trainindex_x train_10p_index.csv --trainindex_u train_90p_index.csv \
  --dist-url 'tcp://localhost:10001' --multiprocessing-distributed --world-size 1 --rank 0 \
  /path/to/imagenet

Linear Classification and Finetuning

With a pre-trained model, to train a supervised linear classifier on frozen features/weights (e.g. MoCo) on 10% imagenet, run:

python main_lincls.py \
  -a resnet50 \
  --lr 30.0 \
  --epochs 50 --schedule 30 40 \
  --eval-freq 5 \
  --trainindex train_10p_index.csv \
  --model-prefix encoder_q \
  --pretrained /path/to/model_best.pth.tar \
  --dist-url 'tcp://localhost:10001' --multiprocessing-distributed --world-size 1 --rank 0 \
  /path/to/imagenet

To finetune the self-supervised pretrained model on 10% imagenet, with different learning rates for pretrained backbone and last classification layer, run:

python main_cls.py \
  -a resnet50 \
  --lr 0.001 --lr-classifier 0.1 \
  --epochs 50 --schedule 30 40 \
  --eval-freq 5 \
  --trainindex train_10p_index.csv \
  --model-prefix encoder_q \
  --self-pretrained /path/to/model_best.pth.tar \
  --dist-url 'tcp://localhost:10001' --multiprocessing-distributed --world-size 1 --rank 0 \
  /path/to/imagenet

For BYOL, change to --model-prefix online_net.backbone. For the best performance, follow the learning rate setting in Section 5.2 in the paper.

Models

Our pre-trained ResNet-50 models can be downloaded as following:

name	epoch	acc@1% IN	acc@10% IN	acc@100% IN	model
MoCo-EMAN	200	48.9	60.5	67.7	download
MoCo-EMAN	800	55.4	64.0	70.1	download
MoCo-2X-EMAN	200	56.8	65.7	72.3	download
BYOL-EMAN	200	55.1	66.7	72.2	download

License

This project is under the CC-BY-NC 4.0 license. See LICENSE for details.