hust-lidelong / mae Goto Github PK

This is a PyTorch/GPU re-implementation of the paper Masked Autoencoders Are Scalable Vision Learners:

@Article{MaskedAutoencoders2021,
  author  = {Kaiming He and Xinlei Chen and Saining Xie and Yanghao Li and Piotr Doll{\'a}r and Ross Girshick},
  journal = {arXiv:2111.06377},
  title   = {Masked Autoencoders Are Scalable Vision Learners},
  year    = {2021},
}

• The original implementation was in TensorFlow+TPU. This re-implementation is in PyTorch+GPU.

• This repo is a modification on the DeiT repo. Installation and preparation follow that repo.

• This repo is based on timm==0.3.2, for which a fix is needed to work with PyTorch 1.8.1+.

• Visualization demo
• Pre-trained checkpoints + fine-tuning code
• Pre-training code

Run our interactive visualization demo using Colab notebook (no GPU needed):

The following table provides the pre-trained checkpoints used in the paper, converted from TF/TPU to PT/GPU:

The fine-tuning instruction is in FINETUNE.md.

By fine-tuning these pre-trained models, we rank #1 in these classification tasks (detailed in the paper):

The pre-training instruction is in PRETRAIN.md.

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