DailyMAE: Towards Pretraining Masked Autoencoders in One Day

📄 arXiv: 2404.00509v1 📥 PDF

作者: Jiantao Wu, Shentong Mo, Sara Atito, Zhenhua Feng, Josef Kittler, Muhammad Awais

分类: cs.LG, cs.CV

发布日期: 2024-03-31

🔗 代码/项目: GITHUB

💡 一句话要点

提出高效的Masked Autoencoders预训练方法以解决计算资源不足问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 自监督学习 掩蔽图像建模 高效训练 计算机视觉 深度学习

📋 核心要点

  1. 现有的自监督学习方法在预训练过程中面临高计算需求，尤其在学术环境中，资源有限导致研究进展缓慢。
  2. 本文提出了一种高效的训练方案，通过优化数据加载和采用渐进式训练技术，旨在提高Masked Autoencoders的训练效率。
  3. 实验结果表明，使用该方法在ImageNet 1K数据集上训练MAE-Base/16模型的速度提升达5.8倍，显著降低了训练时间。

📝 摘要（中文）

近年来，掩蔽图像建模（MIM）作为一种重要的自监督学习（SSL）方法，因其在无标签数据中学习数据表示的有效性而受到关注。然而，预训练的高计算需求在学术环境中造成了显著挑战，阻碍了SSL研究的进展。本文提出了一种高效的MIM基于SSL的训练方案，旨在缓解数据加载瓶颈，并采用渐进式训练技术等方法，以保持预训练性能。通过在单台配备8个A100 GPU的机器上，使用ImageNet 1K数据集训练MAE-Base/16模型800个epoch，仅需18小时，速度提升达5.8倍。这项工作不仅展示了高效SSL训练的可行性，还为SSL研究的更广泛可及性和进步铺平了道路。

🔬 方法详解

问题定义：本文旨在解决自监督学习中Masked Autoencoders预训练的高计算需求问题，现有方法在资源有限的环境中难以实施。

核心思路：通过优化数据加载过程和采用渐进式训练等技术，提升训练效率，同时保持模型的预训练性能。

技术框架：整体架构包括数据加载模块、训练过程优化模块和模型评估模块。数据加载模块通过并行处理加速数据读取，训练过程优化模块则采用渐进式训练策略。

关键创新：最重要的创新在于提出了一种高效的训练配方，使得在单台机器上进行大规模预训练成为可能，显著降低了时间成本。

关键设计：在参数设置上，采用了适合大规模训练的超参数配置，损失函数选择了适合掩蔽学习的设计，网络结构则基于现有的MAE架构进行优化。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，使用本文提出的方法在ImageNet 1K数据集上训练MAE-Base/16模型的速度提升达5.8倍，训练时间缩短至18小时，显著提高了训练效率，为自监督学习的研究提供了新的可能性。

🎯 应用场景

该研究的潜在应用领域包括计算机视觉中的图像分类、目标检测和图像生成等任务。通过降低自监督学习的计算成本，研究者和开发者能够更容易地进行原型设计和初步测试，从而推动相关技术的进步与应用。

📄 摘要（原文）

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.