无论您专注于训练还是推理,数据加载都是深度学习工作流程的一个关键方面。然而,它通常会带来一个矛盾:需要同时具备高度便捷和可定制的解决方案。这两个目标众所周知很难协调。
此问题的传统解决方案之一是扩展处理并并行化用户编写的函数。在这种方法中,用户创建自定义算法,而系统则负责在同时计算任务的多个工作进程中扩展其执行。这就是 torch.DataLoader
发挥作用的地方。
本文记录了我们通过从进程切换到线程来优化 torch.DataLoader
的实验。这项探索之所以成为可能,是因为 Python 不断努力删除 GIL,使我们能够重新思考深度学习工作流程中的并行性,并探索新的性能优化。
什么是 torch.DataLoader?工作原理是什么?
torch.DataLoader
是 PyTorch 中的基础工具,有助于在深度学习应用中加载数据。它在管理数据输入模型的方式方面发挥着关键作用,可确保流程高效且有效。
torch.DataLoader
的重要特性是,它能够并行化加载过程,这在处理大型数据集时至关重要。
这种并行化通常通过创建多个工作进程来实现,每个进程负责加载部分数据,这些进程并行运行,从而能够在训练模型的同时加载和预处理数据。
并行性对于保持稳定的GPU数据流、尽量减少空闲时间和尽量提高资源利用率尤为重要。
可怕的 GIL
torch.DataLoader
使用进程来并行化数据加载任务,这种方法直接源于 Python 架构的一个基本方面,即全局解释器锁(GIL)。
GIL 是一种互斥体,可防止多个原生线程在 CPython (最广泛使用的 Python 实现) 中同时执行 Python 字节码。这锁的引入目的是简化内存管理,并通过在多个线程试图同时访问或修改 Python 对象时防止出现竞争条件,以确保线程安全。
虽然 GIL 使 Python 的内存管理变得简单,并有助于避免复杂的并发错误,但它也施加了一个重大限制:Python 线程并非真正的并行。
在受 CPU 限制的任务中,处理能力是瓶颈,线程不得不轮流运行,导致性能不佳。这就是为什么 torch.DataLoader
使用进程而不是线程的原因。每个进程都在自己的内存空间中运行,完全绕过 GIL,并允许在多核处理器上真正并行执行。
当然,GIL 的影响并非完全是负面的。它通过减少开发者对线程安全的关注来简化 Python 程序的开发,这也是 Python 如此受欢迎的原因之一。
另一方面,GIL 可能会成为 CPU 受限和多线程应用程序的瓶颈,因为它阻碍了多核系统的充分利用。这种权衡在 Python 社区中引发了关于其优缺点的持续争论。
线程交换进程
随着近期的发展,GIL 将在即将推出的 Python 版本中删除。这为 Python 应用程序(包括深度学习)中的并行性开辟了新的可能性。
我们的一个关键想法是尝试将 torch.DataLoader
中基于进程的并行与基于线程的并行交换(图 1)。
使用线程代替进程有几个潜在优势。线程通常比进程轻,从而加快上下文切换速度,并降低内存开销。
然而,线程也带来了它自己的挑战,特别是在确保线程安全和避免死锁等问题方面。
我们实施了基于线程的 torch.DataLoader
版本来探索这些可能性。结果很有趣,并且表明,在某些情况下,线程可以成为进程的可行替代方案。
基于线程的数据加载结果
为了评估在 torch.DataLoader
中使用线程替换进程对性能的影响,我们在不同的数据处理场景中进行了一系列实验。结果突出了基于线程的并行性的潜力和局限性。
使用 nvImageCodec 进行图像解码
在使用 nvImageCodec 的图像解码场景中,出现了使用线程的最引人注目的案例之一。在这种场景中,与传统的基于进程的方法相比,使用线程可显著提高速度。
基准测试详细信息:EPYC 9654 | H100 | 批量大小:512 | 图像大小:640 x 408(JPEG)
实现这一改进的主要原因是 CUDA 上下文切换的减少。在切换上下文时,进程会引入更高的开销,这会导致严重的延迟,尤其是在 GPU 加速的工作负载中。
另一方面,线程可以减少这种开销,从而实现更快、更高效的执行。
使用 Pillow 进行图像解码
与 nvImageCodec
不同,我们对 Pillow (一种广泛使用的 Python 图像库) 的实验表明,线程方法比基于进程的方法略慢。
基准测试详细信息:EPYC 9654 | 批量大小:512 | 图像大小:640 x 408(JPEG)
这里的关键区别在于全局状态的管理方式。Pillow 的操作涉及对存储在字典中的全局状态数据的频繁访问。当多个线程同时访问这些共享数据时,当前的实现依赖于原子来安全地管理这些操作。
然而,atomics 可能会成为争用的瓶颈,与每个工作者都有自己的隔离状态的单独进程相比,这会导致性能降低。
由于这一瓶颈,我们在 discuss.python.org 上发起了一场讨论,重新探讨冻结数据类型的想法,这可以通过实现更高效的读取访问而无需昂贵的原子来帮助缓解这些性能问题。
综合结果:nvImageCodec 与 Pillow
为了更好地显示性能差异,我们将 nvImageCodec
和 Pillow 场景的结果合并到一张图表中(图 4)。
基准测试详细信息:EPYC 9654 | H100 | 批量大小:512 | 图像大小:640 x 408(JPEG)
这种比较清楚地表明了两种方法之间的鲜明对比:
nvImageCodec
:线程的性能明显优于进程,这表明在依赖 CUDA 的 GPU 密集型任务中,线程方法非常有利。- Pillow:进程仍然保持着微小的优势,这表明涉及共享状态的任务可能无法从线程中获益。
这些发现强调,在基于 GPU 的场景中,移除 GIL 可以立即显著提高速度。然而,随着 Python 迈出了进入自由线程领域的第一步,我们应该更加努力地引入新的工具和概念,充分利用硬件功能并充分发挥语言的潜力。
基于线程的 Torch.DataLoader 的优缺点
虽然基于线程的 torch.DataLoader
在某些情况下表现出明显优势,但务必要权衡利弊。
优势显而易见:
- 开销更低:线程的资源密集程度低于进程,因此内存占用率更低,上下文切换速度更快。
- 在某些情况下提供更好的性能:正如
nvImageCodec
实验所示,线程可以减少同步开销,从而提高整体性能。
缺点如下:
- 线程安全:确保代码的线程安全可能具有挑战性,尤其是在复杂的数据管道中。对于线程,总是存在更高的死锁风险,这可能会中断整个数据加载过程。
- 广泛的同步:通常情况下,线程必须比进程更频繁地同步。实现基于线程的执行需要在开发过程中进行更多的审查。
- 迁移现有实现:自由线程的 Python 生态系统尚处于开发的早期阶段。调整深度学习项目具有的大量依赖项需要一些时间。
结束语
删除 GIL 为优化 Python 中的深度学习工作流程带来了新的机会。我们对基于线程的 torch.DataLoader
的探索表明,每当工作者实现涉及 GPU 处理时,它都是一个有益的方法。
然而,对于 CPU 操作而言,由于对数据结构的并行读取访问效率低下,性能往往会出现瓶颈,我们希望在未来能够解决这一问题。
随着 Python 的不断发展,深度学习中数据加载的格局必将发生变化,我们很高兴能走在这些发展的前沿。
如果您有兴趣详细了解我们使用自由线程 Python 进行的实验,请参阅我们的自由线程 Docker 环境。不妨在“issues”部分中发布您的问题,并在您的用例中试用自由线程 Python!