多线程优化数据加载效率
多线程优化数据加载效率
随着Python不断努力删除GIL（全局解释器锁），这为深度学习工作流程中的并行性开辟了新的可能性。本文记录了NVIDIA通过从进程切换到线程来优化PyTorch的数据加载器（torch.DataLoader）的实验，详细介绍了GIL对Python多线程的影响，并通过实验证明了在某些场景下使用线程可以显著提高数据加载效率。
无论您专注于训练还是推理，数据加载都是深度学习工作流程的一个关键方面。然而，它通常会带来一个矛盾：需要同时具备高度便捷和可定制的解决方案。这两个目标众所周知很难协调。
此问题的传统解决方案之一是扩展处理并并行化用户编写的函数。在这种方法中，用户创建自定义算法，而系统则负责在同时计算任务的多个工作进程中扩展其执行。这就是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）。
a) 基于进程的并行性
b) 基于线程的并行性图 1 并行度比较
使用线程代替进程有几个潜在优势。线程通常比进程轻，从而加快上下文切换速度，并降低内存开销。然而，线程也带来了它自己的挑战，特别是在确保线程安全和避免死锁等问题方面。
我们实施了基于线程的torch.DataLoader版本来探索这些可能性。结果很有趣，并且表明，在某些情况下，线程可以成为进程的可行替代方案。
基于线程的数据加载结果
为了评估在torch.DataLoader中使用线程替换进程对性能的影响，我们在不同的数据处理场景中进行了一系列实验。结果突出了基于线程的并行性的潜力和局限性。
使用 nvImageCodec 进行图像解码
在使用nvImageCodec的图像解码场景中，出现了使用线程的最引人注目的案例之一。在这种场景中，与传统的基于进程的方法相比，使用线程可显著提高速度。
图 2.nvImageCodec在两种场景下的吞吐量，越高越好。
基准测试详细信息：EPYC 9654 | H100 | 批量大小：512 | 图像大小：640 x 408（JPEG）
实现这一改进的主要原因是CUDA上下文切换的减少。在切换上下文时，进程会引入更高的开销，这会导致严重的延迟，尤其是在GPU加速的工作负载中。另一方面，线程可以减少这种开销，从而实现更快、更高效的执行。
使用 Pillow 进行图像解码
与nvImageCodec不同，我们对Pillow（一种广泛使用的Python图像库）的实验表明，线程方法比基于进程的方法略慢。
图 3：在两种情况下，Pillow的吞吐量越高越好。
基准测试详细信息：EPYC 9654 | 批量大小：512 | 图像大小：640 x 408（JPEG）
这里的关键区别在于全局状态的管理方式。Pillow的操作涉及对存储在字典中的全局状态数据的频繁访问。当多个线程同时访问这些共享数据时，当前的实现依赖于原子来安全地管理这些操作。然而，atomics可能会成为争用的瓶颈，与每个工作者都有自己的隔离状态的单独进程相比，这会导致性能降低。
由于这一瓶颈，我们在discuss.python.org上发起了一场讨论，重新探讨冻结数据类型的想法，这可以通过实现更高效的读取访问而无需昂贵的原子来帮助缓解这些性能问题。
综合结果：nvImageCodec 与 Pillow
为了更好地显示性能差异，我们将nvImageCodec和Pillow场景的结果合并到一张图表中（图 4）。
图 4. Pillow 和nvImageCodec 使用基于线程和进程的torch.DataLoader 的组合吞吐量。
基准测试详细信息：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的不断发展，深度学习中数据加载的格局必将发生变化，我们很高兴能走在这些发展的前沿。