简单粗暴！Meta公开加速ViT训练流程，吞吐量提升4倍！

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

点击进入—> CVer 微信技术交流群

转载自：机器之心

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

现阶段，视觉 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/

点击进入—> CV 微信技术交流群

ICCV和CVPR 2021论文和代码下载

后台回复：CVPR2021，即可下载CVPR 2021论文和代码开源的论文合集
后台回复：ICCV2021，即可下载ICCV 2021论文和代码开源的论文合集
后台回复：Transformer综述，即可下载最新的3篇Transformer综述PDF

目标检测和Transformer交流群成立

扫描下方二维码，或者添加微信：CVer6666，即可添加CVer小助手微信，便可申请加入CVer-目标检测或者Transformer 微信交流群。另外其他垂直方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer等。

一定要备注：研究方向+地点+学校/公司+昵称（如目标检测或者Transformer+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲扫码或加微信: CVer6666，进交流群

CVer学术交流群（知识星球）来了！想要了解最新最快最好的CV/DL/ML论文速递、优质开源项目、学习教程和实战训练等资料，欢迎扫描下方二维码，加入CVer学术交流群，已汇集数千人！

▲扫码进群
▲点击上方卡片，关注CVer公众号
整理不易，请点赞和在看