MAE — MMPretrain 1.2.0 documentation (original) (raw)
Abstract¶
This paper shows that masked autoencoders (MAE) are scalable self-supervised learners for computer vision. Our MAE approach is simple: we mask random patches of the input image and reconstruct the missing pixels. It is based on two core designs. First, we develop an asymmetric encoder-decoder architecture, with an encoder that operates only on the visible subset of patches (without mask tokens), along with a lightweight decoder that reconstructs the original image from the latent representation and mask tokens. Second, we find that masking a high proportion of the input image, e.g., 75%, yields a nontrivial and meaningful self-supervisory task. Coupling these two designs enables us to train large models efficiently and effectively: we accelerate training (by 3× or more) and improve accuracy. Our scalable approach allows for learning high-capacity models that generalize well: e.g., a vanilla ViT-Huge model achieves the best accuracy (87.8%) among methods that use only ImageNet-1K data. Transfer performance in downstream tasks outperforms supervised pretraining and shows promising scaling behavior.
How to use it?¶
from mmpretrain import inference_model
predict = inference_model('vit-base-p16_mae-300e-pre_8xb128-coslr-100e_in1k', 'demo/bird.JPEG') print(predict['pred_class']) print(predict['pred_score'])
Models and results¶
Pretrained models¶
Model | Params (M) | Flops (G) | Config | Download |
---|---|---|---|---|
mae_vit-base-p16_8xb512-amp-coslr-300e_in1k | 111.91 | 17.58 | config | model | log |
mae_vit-base-p16_8xb512-amp-coslr-400e_in1k | 111.91 | 17.58 | config | model | log |
mae_vit-base-p16_8xb512-amp-coslr-800e_in1k | 111.91 | 17.58 | config | model | log |
mae_vit-base-p16_8xb512-amp-coslr-1600e_in1k | 111.91 | 17.58 | config | model | log |
mae_vit-large-p16_8xb512-amp-coslr-400e_in1k | 329.54 | 61.60 | config | model | log |
mae_vit-large-p16_8xb512-amp-coslr-800e_in1k | 329.54 | 61.60 | config | model | log |
mae_vit-large-p16_8xb512-amp-coslr-1600e_in1k | 329.54 | 61.60 | config | model | log |
mae_vit-huge-p16_8xb512-amp-coslr-1600e_in1k | 657.07 | 167.40 | config | model | log |
Image Classification on ImageNet-1k¶
Model | Pretrain | Params (M) | Flops (G) | Top-1 (%) | Config | Download |
---|---|---|---|---|---|---|
vit-base-p16_mae-300e-pre_8xb128-coslr-100e_in1k | MAE 300-Epochs | 86.57 | 17.58 | 83.10 | config | N/A |
vit-base-p16_mae-400e-pre_8xb128-coslr-100e_in1k | MAE 400-Epochs | 86.57 | 17.58 | 83.30 | config | N/A |
vit-base-p16_mae-800e-pre_8xb128-coslr-100e_in1k | MAE 800-Epochs | 86.57 | 17.58 | 83.30 | config | N/A |
vit-base-p16_mae-1600e-pre_8xb128-coslr-100e_in1k | MAE 1600-Epochs | 86.57 | 17.58 | 83.50 | config | model | log |
vit-base-p16_mae-300e-pre_8xb2048-linear-coslr-90e_in1k | MAE 300-Epochs | 86.57 | 17.58 | 60.80 | config | N/A |
vit-base-p16_mae-400e-pre_8xb2048-linear-coslr-90e_in1k | MAE 400-Epochs | 86.57 | 17.58 | 62.50 | config | N/A |
vit-base-p16_mae-800e-pre_8xb2048-linear-coslr-90e_in1k | MAE 800-Epochs | 86.57 | 17.58 | 65.10 | config | N/A |
vit-base-p16_mae-1600e-pre_8xb2048-linear-coslr-90e_in1k | MAE 1600-Epochs | 86.57 | 17.58 | 67.10 | config | N/A |
vit-large-p16_mae-400e-pre_8xb128-coslr-50e_in1k | MAE 400-Epochs | 304.32 | 61.60 | 85.20 | config | N/A |
vit-large-p16_mae-800e-pre_8xb128-coslr-50e_in1k | MAE 800-Epochs | 304.32 | 61.60 | 85.40 | config | N/A |
vit-large-p16_mae-1600e-pre_8xb128-coslr-50e_in1k | MAE 1600-Epochs | 304.32 | 61.60 | 85.70 | config | N/A |
vit-large-p16_mae-400e-pre_8xb2048-linear-coslr-90e_in1k | MAE 400-Epochs | 304.33 | 61.60 | 70.70 | config | N/A |
vit-large-p16_mae-800e-pre_8xb2048-linear-coslr-90e_in1k | MAE 800-Epochs | 304.33 | 61.60 | 73.70 | config | N/A |
vit-large-p16_mae-1600e-pre_8xb2048-linear-coslr-90e_in1k | MAE 1600-Epochs | 304.33 | 61.60 | 75.50 | config | N/A |
vit-huge-p14_mae-1600e-pre_8xb128-coslr-50e_in1k | MAE 1600-Epochs | 632.04 | 167.40 | 86.90 | config | model | log |
vit-huge-p14_mae-1600e-pre_32xb8-coslr-50e_in1k-448px | MAE 1600-Epochs | 633.03 | 732.13 | 87.30 | config | model | log |
Citation¶
@article{He2021MaskedAA, title={Masked Autoencoders Are Scalable Vision Learners}, author={Kaiming He and Xinlei Chen and Saining Xie and Yanghao Li and Piotr Doll'ar and Ross B. Girshick}, journal={arXiv}, year={2021} }