MAE is available in HuggingFace Transformers (both PyTorch and TF)

Hi everyone!

Thanks for this great work. Since several months, MAE is available in HuggingFace Transformers with an easy-to-use and familiar API: https://huggingface.co/docs/transformers/model_doc/vit_mae

The model is available both in PyTorch and TensorFlow.

The API is the same as BertModel for instance (but it's called ViTMAEModel). There's also a ViTMAEForPreTraining class, similar to BertForPreTraining, which includes the decoder and loss calculation.

We do provide an easy-to-use script to directly pre-train a ViTMAE model on your custom data, called run_mae.py. It can be found here: https://github.com/huggingface/transformers/tree/main/examples/pytorch/image-pretraining. The script leverages the Trainer and runs directly on a distributed environment. The script also allows for logging to Tensorboard/Weights and Biases, and pushing the model to the hub during and/or after training (with easy reloading afterwards using from_pretrained).

After pre-training, you can automatically load the weights into a ViTForImageClassification model, allowing for fine-tuning. Fine-tuning notebooks can be found here (for Torchvision for data augmentation) or here (in case you prefer Albumentations for data augmentation).

I also created a notebook to illustrate visualization (based on this repo): https://github.com/NielsRogge/Transformers-Tutorials/blob/master/ViTMAE/ViT_MAE_visualization_demo.ipynb