maskedautoencoder 0.0.1

Masked Autoencoders: An Updated PyTorch Implementation for Single GPU with 4GB Memory

This is an updated PyTorch/GPU re-implementation of the paper Masked Autoencoders Are Scalable Vision Learners for consumer GPU user for learning purpose.

• Updated to latest Torch and Timm
• Use Imagenette as the default dataset so that you can run the training in a consumer GPU to debug the code immediately without downloading the huge Imagenet

Github Repo: 🔗

Command to Train the model:

pip install maskedautoencoder

git checkout https://github.com/henrywoo/mae/
cd mae
pip install -r requirements.txt
bash run.sh

Screenshot of training it with a 4G GPU laptop:

One liner change to replace ImageNette with ImageNet1K:

Repalce

dataset_train = get_cv_dataset(path=DS_PATH_IMAGENETTE, transform=transform_train, name="full_size")

with

dataset_train = get_cv_dataset(path=DS_PATH_IMAGENET1K, transform=transform_train)

Catalog

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

Visualization demo

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

Fine-tuning with pre-trained checkpoints

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):

Pre-training

The pre-training instruction is in PRETRAIN.md.

License

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

Other Versions

• The original version: PyTorch Version
• Other version: TF, MAE-pytorch 1, MAE-pytorch 2