FFCV: Accelerating Training by Removing Data Bottlenecks

FFCV: Accelerating Training by Removing Data Bottlenecks

相关链接：arxiv
关键字：机器学习、数据加载、训练加速、ImageNet、FFCV库

摘要

我们介绍了FFCV，一个用于简单快速机器学习模型训练的库。FFCV通过消除训练过程中的（通常是微妙的）数据瓶颈来加速模型训练。具体来说，我们结合了高效的文件存储格式、缓存、数据预加载、异步数据传输和即时编译等技术，以确保（a）数据加载和传输的效率大大提升，保证GPU能够达到完全利用；以及（b）尽可能将更多的数据处理任务卸载到CPU异步执行，释放GPU周期用于训练。使用FFCV，我们在ImageNet数据集上训练了ResNet-18和ResNet-50，并且在准确度和训练时间之间实现了竞争性的折衷。例如，在单台机器上仅用20分钟我们就能够将ImageNet ResNet-50模型训练到75%的准确率。我们通过几个案例研究展示了FFCV的性能、易用性、可扩展性以及适应资源限制的能力。

核心方法

• 预处理与训练时加载：FFCV首先预处理数据集以便于高吞吐量加载，然后在训练时加载阶段，FFCV的数据加载器将替换原有学习系统的数据加载器，无需进行其他实现更改。
• 改进的文件格式：FFCV引入了新的文件格式.beton，减轻文件系统压力，并提高灵活性和搜索/索引性。
• 高效的数据读取：FFCV提供了多种针对不同计算环境的高吞吐量读取策略，包括操作系统缓存、进程缓存和准随机采样。
• 快速数据处理：FFCV通过即时（JIT）编译数据处理管道来平衡效率与灵活性，自动将用户提供的Python数据处理管道编译成优化的机器码。
• 规避数据传输成本：由于编译后的机器码不受Python解释器的监管，FFCV可以使用线程而不是子进程，并在训练循环中异步并行运行数据准备操作。

实验说明

	FFCV 训练时间	PyTorch 训练时间	提升情况
单模型训练 (1x ResNet-50)	20分钟	大于1小时	远远超过
多模型训练 (8x ResNet-18)	无额外开销	带宽压力增加导致显著减速	显著提升
低内存训练环境	最小化性能损失	严重受限	显著提升
非计算机视觉任务	显著加速（与PyTorch对比）	没有数据	显著提升

• 通过在各种环境（AWS EC2 p4d.24xlarge 机器和使用网络文件系统训练环境）对FFCV和标准PyTorch数据加载器进行对比，结果显示FFCV实现了显著的速度提升。
• FFCV在训练单一模型时，可以在20分钟内训练一个ResNet-50模型到75%的准确率，远远超过了使用PyTorch的训练时间。
• FFCV能够在不增加任何额外开销的情况下同时训练多个模型，这在进行超参数搜索时非常有用。
• 即使在内存有限的设置中，FFCV仍可以通过进程级缓存和准随机加载策略实现快速训练。
• 除了计算机视觉任务，FFCV还可以应用于其他领域，比如使用稀疏线性回归解算器，在仅更换数据加载器的情况下，实现了1.6倍的加速。

结论

我们提出了FFCV，一个优化的框架，用于消除机器学习模型训练过程中的数据瓶颈。我们使用FFCV显著改善了ImageNet数据集的速度/准确度折衷，并通过一系列案例研究展示了FFCV加速几乎任何机器学习工作负载的潜力。我们工作的主要局限性在于它没有解决训练中的非数据瓶颈问题，因此在数据加载和处理较快的设置中（例如，自然语言处理）或模型非常大并主导训练时间的设置中（例如，大型预训练视觉模型），可能会带来较少（但仍然非零）的改进。