作者:
(1) Iason Ofeidis,耶鲁大学电气工程系、耶鲁网络科学研究所,纽黑文{同等贡献};
(2)Diego Kiedanski,耶鲁大学电气工程系、耶鲁网络科学研究所,纽黑文{同等贡献};
(3) Leandros TassiulasLevon Ghukasyan,Activeloop,美国加利福尼亚州山景城,电气工程系,耶鲁大学网络科学研究所,纽黑文。
我们选择了几个库和数据集来比较它们的特性和性能。尽管我们尽力让这些库和数据集尽可能通俗易懂,但数据加载领域仍在不断发展,每天都会添加新的库和版本。在这方面,我们希望以下列表能够很好地概述数据加载器的当前功能,而不一定声称或找到总体上最好的(这可能会在撰写本文和发布时发生变化)。我们将实验的所有源代码都提供给公众,并希望结果完全可重现。
我们选择了七个库来进行实验:PyTorch(Paszke 等人,2019 年)、Torchdata(TorchData,2021 年)、Hub(团队,2022a 年)、FFCV(Leclerc 等人,2022 年)、Webdatasets(Webdataset,2013 年)、Squirrel(团队,2022b 年)和 Deep Lake(Hambardzumyan 等人,2022 年)。我们发现一件有趣的事情是,并非所有库都支持相同的功能。例如,我们无法使用托管在 S3 存储桶中的数据集运行 FFCV 来执行我们的远程实验。正如我们在第 1 节中提到的,我们在 PyTorch 中运行所有实验。我们考虑过在其他流行的机器学习框架中重现实验,但我们决定放弃这个想法,因为第二个候选者是 Tensorflow,但有传言称谷歌正在放弃它而转向 JAX。图 1 描绘了过去 12 个月中不同 ML 框架的流行程度。
关于数据集,我们最初选择了两个经典数据集来支持两个不同的学习任务:用于图像分类的 CIFAR-10(Krizhevsky 等人,2009)和用于物体检测的 CoCo(Lin 等人,2014)。在进行一些实验时,我们观察到了奇怪的行为(库的表现比预期的要好),这可以通过以下方式解释
CIFAR-10 适合内存[2]。为此,我们建立了第三个数据集,名为 RANDOM,由随机生成的 256 x 256 像素彩色图像和 20 个随机类别组成。第三个数据集包含 45000 张用于训练的图像、5000 张用于验证的图像和 500 张用于测试的图像,它比 CIFAR-10 大得多。
我们对所有库使用相同的转换,以使基准具有可比性。唯一的例外是 FFCV,它有自己的不同转换实现。对于图像分类,转换堆栈由以下内容组成:随机水平翻转、规范化、剪切、转换为张量。
对于对象检测,我们主要依赖 Albumentations(Buslaev 等人,2020 年)的转换实现。堆栈如下所示:随机大小裁剪、随机水平翻转、规范化、转换为张量。这些转换适用于图像和边界框。
如果可能,我们会在本地和 S3 等效存储桶中托管数据集。这使我们能够评估通过网络数据流进行训练所导致的速度减慢。我们将在第 4 节中详细介绍训练设置。
鉴于最密集的训练作业涉及使用多个 GPU,因此只要有可能,我们还会在具有多个 GPU 单元的环境中运行相同的实验。由于在运行实验时并非所有库都完全支持 PyTorch Lightning,因此我们决定使用 PyTorch 中的分布式数据并行 (DDP) (Li et al., 2020) 库来实现多 GPU。
对于某些机器学习项目,我们可能只需要访问较大数据集的子集。在这些情况下,能够快速过滤所需的数据点而无需遍历整个数据集可以大大减少总训练时间。一些库允许基于某些特征进行过滤,例如类(用于图像分类任务)。我们探索了使用库提供的过滤方法(如果提供)与根本不过滤相比的速度增益(或损失)。当库不提供过滤方法时,我们简单地实现它们,即扫描整个数据集并仅保留符合指定条件的元素。快速过滤并不一定易于实现,因为它需要维护类似索引的附加结构以避免遍历所有样本。
最后,表 1 列出了不同库与我们在本文中探讨的不同实验和数据集的兼容性。
在进行实验时,我们的主要任务是找到一个客观的指标,以便我们能够以合理的方式比较所有不同的库。
理想的指标应该是训练作业的总运行时间,因为这是我们必须等待和支付的。不幸的是,这将极大地限制我们可以进行的实验数量。经过仔细考虑,我们选择了每秒处理的数据点(图像)的数量,这是我们的数值实验支持的结果。我们考虑了该指标的两种变体:一种是我们使用 ML 模型进行训练并执行反向传播,另一种是我们不使用 ML 模型,只对样本进行迭代,将它们复制到 GPU。这两个指标之间的差异可以从算法 1 中训练循环的伪代码中看出,其中 m 表示速度变量。我们还收集了总运行时间[3]和初始化数据加载器所需的时间。后者的动机是,一些库可能会预先执行昂贵的计算以提高训练时的速度。我们还最终进行了热身以计算速度。这将在第 3.5 小节中进一步讨论。
为了加深对每个库内部机制的理解,我们决定在一次运行中检查执行每个批次以及初始化数据加载器所需的时间。图 3 描述了单个参数组合 [4] 下每个库在算法 1 所述步骤中花费的时间。所有这些实验都涉及 10 个批次后的截止时间。
有趣的是,第一批数据加载器比其他数据加载器花费的时间长得多。这可以解释如下:由于大多数数据加载器此时依赖于延迟加载数据,因此未来的调用将受益于预取、内存中已有的数据和并行化(在 GPU 忙于计算时执行操作)。
1 至 9 号波段的大小可以最好地表明每个文库的扩展情况,因为
大型数据集随该宽度线性增长。我们可以观察到大多数库的宽度均匀,其中 Deep Lake 最短(最快)。另一方面,唯一显示非均匀宽度的库是 FFCV,其中 1 至 3 号波段非常薄,以至于在图像中看不到。
除了 FFCV 和 Deep Lake 之外,所有库的收尾阶段所用时间大致相同,而这两个库的收尾阶段所用时间则要长得多。收尾阶段所用时间主要取决于模型,并不一定能表明每个库的扩展能力如何。
根据此图,我们决定在计算速度时进行预热。这意味着在所有速度计算中忽略第一批所花费的时间。
[2] 在一些评论文献中,这往往是人们所期望的,但我们发现,在涉及小团队和内部工作站的几个实际应用中,情况并非如此。
[3] 这是从模拟开始到截止的时间,实际上通常只有 10 个批次。
[4] 随机数据集,单 GPU,0 个工作器,批量大小 64