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加速几乎任何机器学习工作负载的潜力。我们工作的主要局限性在于它没有解决训练中的非数据瓶颈问题,因此在数据加载和处理较快的设置中(例如,自然语言处理)或模型非常大并主导训练时间的设置中(例如,大型预训练视觉模型),可能会带来较少(但仍然非零)的改进。