MIRA (PyTorch)

Official PyTorch implementation of
Unsupervised Visual Representation Learning via Mutual Information Regularized Assignment
Dong Hoon Lee, Sungik Choi, Hyunwoo Kim, Sae-Young Chung
NeurIPS 2022

Requirements

• torch==1.11.0
• torchvision==0.12.0
• nvidia-dali-cuda110==1.12.0 (optional)
• tqdm
• wandb

We also include environment.yaml for conda environment.

Pretraining

python -m torch.distributed.launch --nproc_per_node=16 train.py --data-dir $DATA_DIR --entity $WANDB_ENTITY --project $WANDB_PROJECT 

Evaluation

Linear evaluation on ImageNet

python -m torch.distributed.launch --nproc_per_node=16 eval_linear.py $MODEL --data-dir $DATA_DIR --entity $WANDB_ENTITY --project $WANDB_PROJECT

Pretrained weights on ImageNet-1k

method	arch	batch-size	epochs	multi-crop	linear eval	download	script
MIRA	RN-50	4096	800	2x224	74.1	model	script
MIRA	RN-50	4096	400	2x224 + 6x96	75.6	model	script

Acknowledgement

Our implementation uses code from the following repositories: DINO, SwAV MoCo-v3, VISSL, and solo-learn

Citation

If you find our work useful, please consider citing it:

@inproceedings{lee2022mira,
    title={Unsupervised Visual Representation Learning via Mutual Information Regularized Assignment},
    author={Lee, Dong Hoon and Choi, Sungik and Kim, Hyunwoo and Chung, Sae-Young},
    booktitle={Advances in Neural Information Processing Systems},
    year={2022}
}