速度提升5.8倍数 | 如果你还在研究MAE或许DailyMAE是你更好的选择，更快更强更节能！！！

集智书童公众号

发布于 2024-04-25 13:38:46

1510

发布于 2024-04-25 13:38:46

最近，作为一种重要的自监督学习方法，掩膜图像建模（MIM）因其能够从无标签数据中学习数据表示的有效性而受到关注。众多研究强调了MIM的优点，突显了在大型数据集上预训练的模型如何提升下游任务的性能。然而，预训练的高计算需求在学术环境中尤其带来了重大挑战，从而阻碍了自监督学习研究的进展。在本研究中，我们提出了针对基于MIM的自监督学习的有效训练方法，重点关注缓解数据加载瓶颈，并采用渐进式训练技巧和其他技巧来保持预训练性能。我们的库能够在配备8个A100 GPU的单台机器上，在18小时内完成对ImageNet 1K数据集上的MAE-Base/16模型进行800轮训练。通过实现最高5.8倍的速度提升，这项工作不仅展示了进行高效自监督训练的可行性，而且为自监督学习研究的广泛可及性和进步铺平了道路，特别是对于自监督学习想法的原型和初步测试。

1 Introduction

自监督学习（SSL）在机器学习中代表了转变性的飞跃，通过利用未标记数据来进行有效的模型训练[3, 4, 20, 22, 31, 32, 33, 34]。这种学习范式得益于大规模数据集，以学习丰富表示用于小样本学习[8]和迁移学习[13, 23]。互联网上大量的未标记数据激发了对深度神经网络模型在大数据集上训练的需求。目前，SSL的成功通常需要在高性能计算集群（HPC）[8, 11, 17]上训练数周。例如，iBOT [47]在16个V100上训练了193小时，用于ViT-S/16。这些计算不包括在开发SSL框架时测试不同假设所需要的时间，这些假设需要在ImageNet-1K[36]的适当规模上进行测试，ImageNet-1K拥有120万个样本，并且需要相当数量的迭代。因此，高效的预训练配方被高度期望以加速SSL算法的研究，例如，超参数调整和新算法的快速验证。为了减少训练时间，一些研究人员在ImageNet-1K[36]的子集上训练他们的模型，例如10%的样本[3]。然而，当模型扩展到大型数据集时，可能会存在性能差距，即在小数据集上表现成熟的模型可能无法处理复杂问题上的多样性。

通过混合精度[30]、FlashAttention[16]、有效架构[22, 38]等方法，人们已经做出了大量努力来提高训练效率。FFCV[27]被提出以消除数据加载的瓶颈，它增加了每秒处理的图像数量。基于此，FFCV-SSL[7]将SimCLR的训练时间从32小时加速到8小时，使用8个V100。Mask自动编码器（MAEs）[22]，通过移除Mask的标记，加速训练（超过3倍或更多）并达到具有竞争力的性能。尽管MAEs的效率很高，但原始实现仍然需要大约800个GPU小时（V100）来复现结果。

在这项工作中，我们旨在优化单台机器上的训练效率。尽管近年来硬件得到了显著改进，例如2022年发布的H100比2017年发布的V100快3倍，但数据加载已经成为训练基础模型的主要瓶颈，如图2所示。然而，大多数机器学习库仍然使用Pytorch实现，从文件夹（IF）中读取图像，这严重阻碍了训练过程。为了消除这个瓶颈，我们利用FFCV[27]进行快速数据加载。我们进一步研究最大分辨率和图像质量对压缩的影响，这对于存储、吞吐量和性能之间的权衡至关重要。此外，我们提出“裁剪解码”以优化FFCV，以实现更快解码和降低内存使用。我们改进的FFCV，称为ESSL，比原始实现快27.6%，节省13.7%的内存。

通过渐进式训练可以进一步加快训练过程，在这个过程中，图像大小在训练过程中逐渐增加[24, 38]。人们已经做出了大量努力在小图像上训练模型。一个优点是减少标记数量，从而降低计算成本。另一个优点是防止过拟合的正则化效果[39, 28]。通过逐渐增强数据增强，模型可以收敛到一个更好的最小值。然而，关于渐进式训练的大多数研究都是基于卷积神经网络（CNN）并关注监督学习。我们重新审视了将渐进式训练应用于ViT以进行微调的想法，因为图像大小与补丁嵌入所看到的明显大小有关。我们进一步将渐进式训练应用于预训练。令人惊讶的是，在训练过程中先逐渐减小后增大图像大小的回文方案，保持了有竞争力的性能并减少了训练时间。我们的贡献可以总结为：

我们提出了一个用于Mask自动编码器的机器学习库，其中我们的mae-base/16预训练过程比官方的快5.8倍。
我们引入“裁剪解码”以减少内存使用并加速加载管道。我们讨论了关于分辨率和图像质量的压缩效果。
我们提出了一种新颖的预训练策略，利用回文方案。这种方法出人意料地保持了有竞争力的性能，同时实现了10.9%的训练时间减少。

2 Related Work

数据加载库。 对于具有单个节点的机器学习系统，由于人工智能芯片的快速发展以加快计算速度，但对数据文件处理方面的改进却很少，因此将数据样本加载到张量中正成为机器学习的瓶颈。一个通用的数据加载流程包括读取和处理。DALI [2] 是 pytorch 数据加载器的性能更高的替代品，它通过使用 GPU 处理数据来提高吞吐量。然而，由于对小文件进行随机读取，从文件夹中读取单个图像会减慢流程。WEBDATASET [1] 通过将样本聚集到 tar 文件中来优化 I/O 速率。通过优化数据加载和处理，Fast Forward Computer Vision (FFCV) [27] 消除了数据瓶颈，结合了诸如高效文件存储格式、缓存、数据预加载、异步数据传输和即时编译等技术。注意到典型的流程会解码整张图像然后裁剪一部分，在我们的工作中，我们引入了 crop decode 操作以避免解码被丢弃的部分，并在 FFCV 中应用它来加速数据加载。

Mask图像建模。 自监督学习正成为机器学习中的一个热门研究课题，以解决监督学习需要大量标记数据的局限性 [12, 20, 34, 42]。Mask图像建模（MIM）是成功的几种方法之一，在从广泛的数据集中学习视觉表示方面表现出色 [4, 44, 22, 43, 46]。它从 NLP [17] 中Mask语言模型的成功中汲取灵感，采用类似的Mask和预测策略，使神经网络能够把握视觉数据的复杂结构和上下文。MIM 在增强下游任务性能方面的有效性，包括图像分类和语义分割，证明了其能够学习可迁移和泛化的特征 [4, 6, 22]。

在 MIM 领域中，已经提出了不同的Mask策略。例如，GMML [4] 集成了组Mask和外星块，有效地破坏和模糊输入数据以增强特征学习。BeiT [6] 使用Mask标记作为Mask块的位置占位符，为处理视觉数据中的遮挡提供了一种新颖的方法。Mask自动编码器 [22] 应用了一种非对称结构，编码器丢弃Mask块，而轻量级解码器从潜在表示和Mask标记重建原始图像。已经提出了许多改进措施来优化效率和有效性。例如，Data2vec 2.0 [5] 通过学习丰富的上下文化目标优化了目标的数据效率。尽管还有其他改进措施 [41, 45, 47] 在性能或成本上表现更高，但我们的工作为了简单起见，关注原始的Mask自动编码器 [22]。我们的框架易于使用，并具有在未来采用这些改进的能力。

渐进式训练。 较低的分辨率（例如192x192）有利于训练过程。一方面，它显著减少了训练时间和内存使用，特别是对于变压器模型。另一方面，它具有正则化效果，迫使模型学习图像的一般感觉 [39, 9]。受到 FixRes [40] 的启发，该研究发现训练集和测试集在处理不同大小对象方面的差异会导致验证性能下降，我们研究了裁剪图像的比例缩放效果。逆缩放法则 [28] 表明，只有大型模型从较少的标记中获益。因此，在分辨率方面，减少泛化误差和训练误差之间存在权衡。渐进式调整大小，即在训练过程中逐渐增加图像大小，已经用于图像分类 [24, 35, 7]，视频摘要 [21]，以及视觉 [26]。这种简单的技术不仅加速了训练过程，还减少了泛化误差。EfficientNetV2 [38] 进一步改进了仅增加图像大小的渐进式训练，通过逐渐添加更强的正则化。然而，这种方法在预训练中的渐进式训练的优点尚未得到适当探索。

3 Efficient Masked Autoencoder

在本研究中，我们开发了一种有效的Mask自编码器训练策略，重点关注预训练和微调阶段。我们主要在ImageNet-1K数据集上采用视觉Transformer基础模型（85M）[18]进行训练。尽管采用其他高效架构，如Swin Transformer [29]和分层模型[25；37]，或先进模型如进化MAE [19]和SdAE [14]，可能会进一步减少训练时间，但我们选择基本的MAE，因其简单性和广泛适用性。这一选择由两个因素驱动：

MAE与各种模型兼容，仅需最小修改；
我们旨在建立一个基础库以促进未来的创新，复杂的模型可能会阻碍新发展的产生。

我们的研究利用增强的FFCV [27]消除数据加载延迟，并采用渐进式训练以逐步调整图像大小而不影响性能。这些简单的方法显著加快了学习过程。图2展示了训练过程的比较，显示了预训练阶段的显著加速。表2展示了我们高效的基准测试。我们区分了官方MAE预训练和微调权重，分别表示为MAE和ft-MAE1，以及我们实现的从零开始的模型，分别称为mae和ft-mae。我们的训练方案总结在表1中。

微调方案。我们的微调方案采用了MAE微调设置，但选择了一种不同的数据增强策略。我们没有使用常见的RandAug [15]，它会自动选择组合增强操作以应用于训练数据，但在数据从压缩中的转变方面存在不足，我们采用了更具抗压缩数据集能力的Three Augmentation（3 Aug）[39]过程。值得注意的是，模型在不同的压缩参数下的验证集上产生了不同的结果。因此，我们将验证集标准化为最大分辨率500和100的质量，表示为

\text{Val}(500\_100)

，增强了模型评估的一致性。这种调整，加上渐进式训练，将8个V100 GPU设置上的微调时间从21小时减少到18小时。

预训练方案。 Mask自编码器[22]，一种流行的自监督架构，通过省略Mask块显著加快了训练速度，实现了3倍的速度提升。然而，由于图1所示的硬件快速进步，数据加载成为了瓶颈。我们的方法通过用改进的FFCV和动态图像分辨率策略替换数据加载器，提高了效率。与传统的渐进式学习方法不同，它逐渐增加分辨率，我们发现对于MAE来说，在训练过程中先降低后增加分辨率更有益。这些改进显著减少了预训练时间，在采用渐进式训练的情况下降至17小时，同时保持了MAE的高性能。

Machine Specification

在我们的论文中，我们在三个平台上进行了实验。具体规格列于表3中。我们采用操作系统缓存来实现内存映射，因此，文件系统不会成为瓶颈。

4 Removing the Data Loading Bottleneck.

在本节中，我们展示了在_eureka_平台上的基准测试（详细内容见附录3.1）。我们探讨了图像压缩与性能之间的权衡。尽管通过降低图像质量或大小来压缩图像可以提高加载吞吐量，但这也会以降低图像保真度为代价，并可能降低整体性能。

我们提出了一种简单而有效的_裁剪解码_策略，在不牺牲图像保真度的前提下，实现了显著的解压缩数据加速和内存使用减少。我们改进版的FFCV [27]，称为ESSL，在90的质量水平和500的最大分辨率下，实现了27.6%的吞吐量提升。

我们进一步探讨了压缩参数对吞吐量、性能和存储需求的影响。为了便于讨论，我们使用表示压缩参数的数据集记法Train/Val(res_quality)。我们得出结论，Train(500_95)在这些因素之间取得了良好的平衡，并被选为训练集的默认设置。Val(500_100)提供了最佳性能，但代价是更高的存储需求。因此，Val(500_100)仅用于报告最终结果。

Crop Decode

随机调整大小裁剪（RandomResizedCrop），这是许多机器学习系统中普遍使用的一种常见数据增强技术，通常涉及在裁剪和调整到所需大小之前解码整张图像（图3(a)）。然而，我们提出了一种新颖的方法，称为“裁剪解码”（crop decode），它只解码目标区域，显著减少了处理开销（图3(b)）。我们将应用了“裁剪解码”的FFCV [27] 称为ESSL。图3(c)比较了使用不同最大图像大小的压缩数据集时，现有FFCV中的RandomResizedCrop实现（90%质量）的数据加载性能。这个基准测试包括从IN1K加载100,000张图像，执行随机调整大小的裁剪至224x224，并应用随机水平翻转。

所有基准测试都是在Eureka平台上执行的，批处理大小为256，64个工作进程，并启用操作系统缓存。通过避免对废弃区域进行不必要的内存分配和解码，我们的“裁剪解码”策略在数据加载上实现了显著的27.6%加速，对于质量为90%，最大图像大小为500的数据集，达到了27,493张图像/秒的吞吐量。这表明ESSL对于包含高质量图像的大型数据集具有潜力。

Compression Parameters for Building an FFCV Dataset

构建一个FFCV数据集需要平衡三个关键因素：吞吐量、存储使用和图像保真度。虽然更高的质量或分辨率可以提高图像保真度，但也会增加存储需求并减慢解码过程。此外，对于大规模数据集来说，存储原始图像是不切实际的（例如，1600分辨率的IN1K将需要788.97 GB）。在这种情况下，硬件的I/O限制成为数据加载的瓶颈。

因此，我们关注压缩数据集，并全面比较相关参数（分辨率和质量）。表4总结了在FFCV中各种压缩设置下的文件大小、吞吐量和性能。需要注意的是，数据集在两个压缩过程中进行处理：一个是在最初以JPEG格式存储时，另一个是在构建FFCV数据集时。此外，JPEG是一种有损压缩算法，意味着即使在最高质量设置（100）下，图像也可能出现一些失真。从结果中，我们有以下关键观察：

最高分辨率和质量设置（500_100）产生最大的文件大小（181 GB），但在准确性方面并不优于默认设置，这表明压缩存在未知的最佳设置。
将质量降低到90会进一步减少文件大小并增加吞吐量，但会牺牲准确性。
将分辨率增加到1000与默认设置相比略微增加了文件大小，但减少了吞吐量，并没有提供显著的准确性好处。

因此，由于在效率（就存储和计算速度而言）和有效性（就模型准确性而言）之间取得了良好的平衡，选择了分辨率为500、质量为95的设置作为默认设置。

Compression Shift

数据偏移指的是两个数据集在统计特性上的差异。虽然压缩过程引入的扭曲对人类肉眼通常是不可察觉的，但它们足以影响机器学习模型的性能。此外，不同实现的JPEG压缩也可能导致数据偏移。

在本研究中，我们比较了两种实现引起的压缩偏移：IF中应用的PIL1和FFCV中应用的jpeg-turbo2。为了评估其影响，我们创建了具有不同压缩质量和图像分辨率的多个验证集。然后，我们使用相同的微调权重（ft-MAE-Base/16）在这些集合上评估分类性能。

表5比较了使用PIL和jpeg-turbo压缩的验证集上的分类性能。有趣的是，我们的结果表明，某些扭曲甚至可以提高验证准确率。jpeg-turbo在相同参数下产生了性能下降。ft-MAE-Base/16在验证集上以100%的质量和500的分辨率保持了高性能。因此，此设置被选为后续评估的默认设置。

用于缓解压缩偏移的三种增强方法。 数据增强是一种常用的技术，通过人为增加训练数据的多样性来解决数据偏移，提高模型的泛化能力。在本研究中，我们比较了两种数据增强策略：RandAug [15] 和三种增强（3 Aug）[39]。我们使用MAE-Base/16初始化模型权重，并在Image Folder（IF）和Train(500_95)数据集上进行微调。然后，我们使用top-1准确率在Val(500_100)和IF数据集上评估其性能。

我们的结果显示，当训练和验证数据之间存在压缩偏移时，使用RandAug训练的模型的准确率显著下降（-0.6%）。相反，3 Aug策略对压缩偏移表现出更大的韧性，只经历了较小的性能下降。此外，表7强调了3 Aug在图像处理效率方面显著优于RandAug。因此，我们将3 Aug纳入到我们的微调配方中，以最小化性能牺牲来缓解压缩偏移。

Discussion

尽管FFCV提供了高效的数据加载，但对数据集构建参数的仔细考虑仍然是至关重要的。我们的实验揭示了由于压缩数据集导致的性能退化，突显了文件大小、图像保真度与模型性能之间的权衡。有趣的是，我们观察到一定程度的可控失真甚至可以增强分类准确性。这表明选择适当的分辨率、质量与压缩算法组合有可能进一步优化性能。我们鼓励社区在构建数据集时优先考虑这一平衡，以便在不牺牲性能的情况下进行高效训练。

5 Progressive Training

基于在Fastai [24]和EfficientNetV2 [38]等高效训练方法中成功应用课程学习（curriculum learning）的基础上，我们引入了渐进式训练（progressive training）用于视觉变换器（Vision Transformers，简称ViTs）的微调和预训练，以进一步加快训练过程。与之前主要关注卷积神经网络（CNNs）和分类任务的工作不同，我们的探索在ViTs的背景下应用渐进式训练。据我们所知，虽然先前的研究探讨了ViTs的动态掩码比例 [19]，但在预训练期间对动态分辨率的探讨仍然较少。

预训练可以从低分辨率图像中受益。Li等人 [28] 针对CLIP训练建立了“反向缩放法则”，提出随着模型规模的增大，由于分辨率[39]的规整化效应，训练所需的图像/文本标记更少。这种效应本质上防止了大模型对高分辨率图像中细粒度细节的过拟合。

认识到CNNs和ViTs之间的关键区别，尤其是ViTs中线性块嵌入（linear patch embeddings）不适应空间不变性，我们首先重新审视了感知比例（perceptual ratio）、视在大小（apparent size）和图像分辨率之间的关系。然后，我们探讨了在分辨率变化期间保持感知比例和视在大小的连贯性对微调的影响。我们的发现表明，ViTs能够处理块分布的变化，并在分辨率缩放时优先考虑感知比例的重要性。

与微调中通过增强逐步增加训练难度不同，我们为预训练提出了一种相反的方法：通过降低分辨率或增加掩码比例来实现难度的逐步增加。有趣的是，降低图像分辨率不仅显著提高了训练速度，而且出人意料地带来了性能的改善。

Perceptual Ratio, Apparent Size, and Resolution

我们考虑以下方式建模一个对象：实际大小的对象在图像中以图像大小

H\times H

被捕获。我们可以通过将图像大小除以

S=R/H

来推断对象的视觉比例，其中

R\times R

是对象的视觉大小。在这个背景下，

R\times R

表示图像中对象的视觉大小。然而，重要的是要承认数据增强技术可以改变图像大小和对象的视觉大小。如文献 [40] 所讨论，在从训练集迁移到验证集时，视觉大小在泛化方面发挥着重要作用。在这里，我们重新审视这一点，并解释感知比例、视觉大小和分辨率之间的关系。感知比例指的是模型可以接收到的图像平均百分比。而视觉大小指的是图像在像素中的实际大小。

我们的讨论基于 RandomResizedCrop（RRC），它由一个比例参数

\sigma

控制，使得

\sigma\sim U([\sigma_{-}^{2},\sigma_{+}^{2}])

和一个宽高比

\alpha

，使得

\ln\alpha\sim U([\ln\alpha_{-},\ln\alpha_{+}])

。RRC 首先以随机的比例和宽高比（为了简化，假设是正方形图像，

\alpha=1

）裁剪一个随机区域。裁剪后，对象的感知比例，反映了最具语义信息的内容，降低到

\sigma

。然而，视觉大小，表示图像中对象的感知大小，保持不变。最后，将裁剪的区域调整到输出大小（h）进一步按

\frac{h}{\sigma H}

缩放视觉大小：

R_{\texttt{apparent}}=\mathbb{E}[R\frac{h}{\sigma H}]=\mathbb{E}[Ch/\sigma]=Ch \frac{3}{2}\frac{\sigma_{+}^{2}-\sigma_{-}^{2}}{\sigma_{+}^{3}-\sigma_{-}^{3}}, \tag{1}

其中

C=R/H

是一个常数（为了简化，

C=1

）。调整大小不会改变对象的比例，因此，感知比例的期望是

R_{\texttt{ratio}}=\mathbb{E}[\sigma]=\frac{2}{3}\frac{\sigma_{+}^{3}-\sigma _{-}^{3}}{\sigma_{+}^{2}-\sigma_{-}^{2}}. \tag{2}

如图4 所示，当生成低分辨率图像时，感知比例更好地保留了狐狸的语义信息，而视觉大小保持了输入块的分部。具有相同视觉大小的低分辨率图像可以被视为局部裁剪 [10]，这显著提升了对比学习的性能。理解这些区别对于优化 ViTs 中的训练效率和性能至关重要。

Finetuning with Dynamic Resolution Scaling

在这一部分，我们在Train(500_95)上对预训练的MAE-Base/16模型进行微调，并在Val(500_100)上评估其性能。微调方法基于表1中的方法。我们实施了涉及三个阶段的动态调整方案，分别持续30、30、40个周期。表8总结了动态调整方案的具体细节。

感知比例。默认方案（方案1） 优先保持感知比例，同时逐渐将图像分辨率从160提高到224。此外，在最后阶段加强了数据增强。方案1、1-和1+探讨了不同增强策略的影响。值得注意的是，“3 Aug+”表示三种增强[39]和ColorJitter (0.3)的组合。

视在大小。方案3 专注于在训练过程中保持视在大小，同时逐步提高感知比例，模仿从局部到全局裁剪的转变。

表9比较了不同渐进式方案下的微调性能。结果显示，模型在初期应用强增强时学习困难。逐渐增强的增强策略导致了更好的泛化效果。

Pretraining with Dynamic Resolution Scaling

尽管数据增强在渐进式训练的微调阶段被证明是有益的，但研究表明，在使用Mask图像建模（MIM）[9, 22]的预训练期间，它可能是有害的。因此，我们探索了两种替代策略,直观地说，减小图像尺寸会增加每个补丁内的Mask区域，因为每个较小的补丁代表了较大比例的对象。这导致Mask的感知比更高。

我们设计了方案3，它以相反的方向逐步改变图像尺寸。与方案2相比，方案3实现了略好的性能，并节省了大约12.62%的训练时间。

其他方案。方案5是为了解决低分辨率图像中的信息丢失问题而提出的。方案6是一种在训练过程中增加Mask比和Mask感知比的超设计。

从结果中可以看出，方案5在微调性能上表现最佳，而方案4在线上探测性能上表现最高。然而，方案6在两项评估中的表现都较差。总之，在较低分辨率（较高的Mask感知比）上训练模型可以增强能力，但由于信息丢失的限制，存在局限性。在利用较低分辨率和减轻信息丢失之间找到正确的平衡是至关重要的。

图像文件夹。我们进一步研究了没有FFCV的情况下渐进式训练的有效性。模型在IF上进行预训练。与使用FFCV数据集的发现相似，表12的结果显示出一致的趋势。方案4，在预训练期间直接增加Mask比，再次展现出在线探测性能的最高提升。此外，与固定尺寸方案相比，它实现了更大的训练时间减少，达到了18.5%。