DailyMAE: Towards Pretraining Masked Autoencoders in One Day

TL;DR

This paper introduces efficient training methods for masked autoencoders in self-supervised learning, enabling high-performance pretraining on ImageNet within 18 hours using accessible hardware, thus making SSL research more feasible and faster.

Contribution

The authors present optimized training recipes and techniques that significantly reduce pretraining time for masked autoencoders without sacrificing performance.

Findings

Achieved 5.8x speedup in pretraining time

Trained MAE-Base/16 on ImageNet 1K in 18 hours

Enabled high-efficiency SSL training on a single machine

Abstract

Recently, masked image modeling (MIM), an important self-supervised learning (SSL) method, has drawn attention for its effectiveness in learning data representation from unlabeled data. Numerous studies underscore the advantages of MIM, highlighting how models pretrained on extensive datasets can enhance the performance of downstream tasks. However, the high computational demands of pretraining pose significant challenges, particularly within academic environments, thereby impeding the SSL research progress. In this study, we propose efficient training recipes for MIM based SSL that focuses on mitigating data loading bottlenecks and employing progressive training techniques and other tricks to closely maintain pretraining performance. Our library enables the training of a MAE-Base/16 model on the ImageNet 1K dataset for 800 epochs within just 18 hours, using a single machine equipped with 8 A100 GPUs. By achieving speed gains of up to 5.8 times, this work not only demonstrates the feasibility of conducting high-efficiency SSL training but also paves the way for broader accessibility and promotes advancement in SSL research particularly for prototyping and initial testing of SSL ideas. The code is available in https://github.com/erow/FastSSL.