Pytorch implementation of MixMAE (CVPR 2023)

This repo is the offcial implementation of the paper MixMAE: Mixed and Masked Autoencoder for Efficient Pretraining of Hierarchical Vision Transformers

@article{MixMAE,
  author  = {Jihao Liu, Xin Huang, Jinliang Zheng, Yu Liu, Hongsheng Li},
  journal = {arXiv:2205.13137},
  title   = {MixMAE: Mixed and Masked Autoencoder for Efficient Pretraining of Hierarchical Vision Transformers},
  year    = {2022},
}

Availble pretrained models

Models	Params (M)	FLOPs (G)	Pretrain Epochs	Top-1 Acc.	Pretrain_ckpt	Finetune_ckpt
Swin-B/W14	88	16.3	600	85.1	base_600ep	base_600ep_ft
Swin-B/W16-384x384	89.6	52.6	600	86.3	base_600ep	base_600ep_ft_384x384
Swin-L/W14	197	35.9	600	85.9	large_600ep	large_600ep_ft
Swin-L/W16-384x384	199	112	600	86.9	large_600ep	large_600ep_ft_384x384

Training and evaluation

We use Slurm for multi-node distributed pretraining and finetuning.

Pretrain

sh exp/base_600ep/pretrain.sh partition 16 /path/to/imagenet

• Training with 16 GPUs on your partition.
• Batch size is 128 * 16 = 2048.
• Default setting is to train for 600 epochs with mask ratio of 0.5.

Finetune

sh exp/base_600ep/finetune.sh partition 8 /path/to/imagenet

• Training with 8 GPUs on your partition.
• Batch size is 128 * 8 = 1024.
• Default setting is to finetune for 100 epochs.