2022/05/27 16:28

不堆参数、不靠时长，Meta加速ViT训练流程，吞吐量4倍提升

有时候，好的训练「技巧」比蛮力堆参更有效。

现阶段，视觉 transformer（ViT）模型已经在图像分类、目标检测与分割等各样各样的计算机视觉任务中得到了广泛应用，并可以在视觉表征与识别中实现 SOTA 结果。由于计算机视觉模型的性能往往与参数量和训练时长呈正相关，AI 社区已经实验了越来越大规模的 ViT 模型。

但应看到，随着模型开始超出万亿次浮点运算的规模，该领域已经遇到了一些主要的瓶颈。训练单个模型可能耗费数月，需要数以千块的 GPU，进而增加了加速器需求并导致大规模 ViT 模型将很多从业者「排除在外」。

为了扩展 ViT 模型的使用范围，Meta AI 的研究者已经开发出了更高效的训练方法。非常重要的一点是对训练进行优化以实现最佳的加速器利用。但是，这一过程耗时费力且需要大量的专业知识。为了设置有序的实验，研究者必须从无数可能的优化方案中进行选择：一次训练过程中执行的百万次运算中的任何一个都有可能受到低效率的影响和阻碍。

Meta AI 发现，通过将一系列优化应用到其图像分类代码库 PyCls 中的 ViT 实现，可以提升计算和存储效率。对于使用 PyCIs 训练的 ViT 模型，Meta AI 的方法可以提升训练速度和每加速器吞吐量（TFLOPS）。

下图展示了使用优化代码库 PyCIs 后每芯片（per chip）加速器吞吐量相较于 V100 基准的相对增加，而 A100 优化的加速器吞吐量是 V100 基准的 4.05 倍。

运行原理

Meta AI 首先对 PyCIs 代码库进行分析以确认低训练效率的潜在来源，最终将注意力放在了对数字格式的选择上。在默认情况下，大多数应用使用 32-bit 单精度浮点格式来表征神经网络值。转换至 16-bit 半精度格式（FP16）可以减少模型的内存占用和执行时间，但往往也会降低准确率。

研究者采取了折中方案，即混合精度。利用它，系统通过单精度格式执行计算以加速训练并减少内存使用，同时通过单精度存储结果以保持准确率。他们没有手动地将部分网络转换至半精度，而是实验了不同模式的自动混合精度训练，这样可以在数字格式之间自动切换。更高级模式的自动混合精度主要依赖半精度运算和模型权重。研究者采用的平衡设置既能大幅度加速训练，同时也不牺牲准确率。

为了使流程更加高效，研究者充分利用了 FairScale 库中的完全分片数据并行（Fully Sharder Data Parallel, FSDP）训练算法，它在 GPU 上对参数、梯度和优化器状态进行分片。通过 FSDP 算法，研究者可以使用更少的 GPU 构建更大量级的模型。此外，研究者还使用了 MTA 优化器、一个池化的 ViT 分类器和一个 batch-second 输入张量布局来跳过冗余转置运算。

下图 X 轴为可能的优化，Y 轴为采用 ViT-H/16 训练时加速器吞吐量相较于分布式数据并行（DDP）基准的相对增加。

研究者在总 patch 大小为 560 时实现了 1.51 倍的加速器吞吐量提升，以每个加速器芯片上每秒执行的浮点运算数量衡量。通过将图像大小从 224 像素增加至 256 像素，他们可以将吞吐量提升至 1.86 倍。但是，改变图像大小意味着超参数的变化，这会对模型的准确率造成影响。在完全 FP16 模式下训练时，相对吞吐量增加至 2.18 倍。尽管有时会降低准确率，但在实验中准确率降低少于 10%。

下图 Y 轴为 epoch 时间，在整个 ImageNet-1K 数据集上一次训练的持续时间。这里专注于现有配置的实际训练时间，这些配置通常使用 224 像素的图像大小。

Meta AI 的研究者使用优化方案，将 epoch 时间（在整个 ImageNet-1K 数据集上一次训练的持续时间）从 0.65 小时减少到 0.43 小时。

下图 X 轴表示特定配置中 A100 GPU 加速器芯片的数量，Y 轴表示每芯片 TFLOPS 的绝对吞吐量。

该研究还讨论了不同 GPU 配置的影响。在每种情况下，系统都实现了比分布式数据并行（DDP）基线水平更高的吞吐量。随着芯片数量的增加，由于设备间通信的开销，我们可以观察到吞吐量略有下降。然而，即使用 64 块 GPU，Meta 的系统也比 DDP 基准快 1.83 倍。

新研究的意义

将 ViT 训练中可实现的吞吐量翻倍可以有效让训练集群规模翻倍，提高加速器利用率直接减少了 AI 模型的碳排放。由于最近大模型的发展带来了更大模型和更长训练时间的趋势，这种优化有望帮助研究领域进一步推动最先进的技术，缩短周转时间并提高生产力。

原文链接：https://ai.facebook.com/blog/significantly-faster-vision-transformer-training/

工程MeTA

相关数据

池化技术

池化（Pooling）是卷积神经网络中的一个重要的概念，它实际上是一种形式的降采样。有多种不同形式的非线性池化函数，而其中“最大池化（Max pooling）”是最为常见的。它是将输入的图像划分为若干个矩形区域，对每个子区域输出最大值。直觉上，这种机制能够有效的原因在于，在发现一个特征之后，它的精确位置远不及它和其他特征的相对位置的关系重要。池化层会不断地减小数据的空间大小，因此参数的数量和计算量也会下降，这在一定程度上也控制了过拟合。通常来说，CNN的卷积层之间都会周期性地插入池化层。

来源：cs231n

权重技术

线性模型中特征的系数，或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0，则相应的特征对模型来说没有任何贡献。

来源：Google AI Glossary

基准技术

一种简单的模型或启发法，用作比较模型效果时的参考点。基准有助于模型开发者针对特定问题量化最低预期效果。

来源：Google ML Glossary

参数技术

在数学和统计学裡，参数（英语：parameter）是使用通用变量来建立函数和变量之间关系（当这种关系很难用方程来阐述时）的一个数量。

来源：维基百科

超参数技术

在机器学习中，超参数是在学习过程开始之前设置其值的参数。相反，其他参数的值是通过训练得出的。不同的模型训练算法需要不同的超参数，一些简单的算法（如普通最小二乘回归）不需要。给定这些超参数，训练算法从数据中学习参数。相同种类的机器学习模型可能需要不同的超参数来适应不同的数据模式，并且必须对其进行调整以便模型能够最优地解决机器学习问题。在实际应用中一般需要对超参数进行优化，以找到一个超参数元组（tuple），由这些超参数元组形成一个最优化模型，该模型可以将在给定的独立数据上预定义的损失函数最小化。

来源：Wikipedia

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数，这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在维空间内，有个分量的一种量，其中每个分量都是坐标的函数，而在坐标变换时，这些分量也依照某些规则作线性变换。称为该张量的秩或阶（与矩阵的秩和阶均无关系）。在数学里，张量是一种几何实体，或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达，记作标量的数组，但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中，表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了，它们都是二阶张量，对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

来源：维基百科

计算机视觉技术

计算机视觉（CV）是指机器感知环境的能力。这一技术类别中的经典任务有图像形成、图像处理、图像提取和图像的三维推理。目标识别和面部识别也是很重要的研究领域。

来源：机器之心

神经网络技术

（人工）神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型，那时候研究者构想了「感知器（perceptron）」的想法。这一领域的研究者通常被称为「联结主义者（Connectionist）」，因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型，它们都是前馈神经网络：卷积神经网络（CNN）和循环神经网络（RNN），其中 RNN 又包含长短期记忆（LSTM）、门控循环单元（GRU）等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习，但也有一些为无监督学习设计的变体，比如自动编码器和生成对抗网络（GAN）。

来源：机器之心

准确率技术

分类模型的正确预测所占的比例。在多类别分类中，准确率的定义为：正确的预测数/样本总数。在二元分类中，准确率的定义为：(真正例数+真负例数)/样本总数

来源：Google ML Glossary

图像分类技术

图像分类，根据各自在图像信息中所反映的不同特征，把不同类别的目标区分开来的图像处理方法。它利用计算机对图像进行定量分析，把图像或图像中的每个像元或区域划归为若干个类别中的某一种，以代替人的视觉判读。

来源：百度百科

优化器技术

优化器基类提供了计算梯度loss的方法，并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法，如梯度下降和Adagrad。优化器是提供了一个可以使用各种优化算法的接口，可以让用户直接调用一些经典的优化算法，如梯度下降法等等。优化器（optimizers）类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类，但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer（tensorflow下的优化器包）等等这些算法。

来源：维基百科

目标检测技术

一般目标检测（generic object detection）的目标是根据大量预定义的类别在自然图像中确定目标实例的位置，这是计算机视觉领域最基本和最有挑战性的问题之一。近些年兴起的深度学习技术是一种可从数据中直接学习特征表示的强大方法，并已经为一般目标检测领域带来了显著的突破性进展。

来源：机器之心