【DataLoader】AI 框架基础技术之数据读取

---

0 概述

数据读取是模型训练的一个基本组成部分。相比于复杂的网络算法设计，数据读取这个概念听起来简单且微不足道。然而，在实际业务落地中，数据读取往往是造成模型速度差和训练精度低的元凶。

不仅如此，在深度学习社区内，数据读取相关话题的讨论热度也一直居高不下，各种各样的工具和解决方案层出不穷——就连英伟达也亲自下场搞了一个数据预处理的库：The NVIDIA Data Loading Library (DALI)。

今天，我们将做一个总述型的介绍，讲讲数据读取中不容小觑的技术点。

1 数据读取基本流程/背景

DataLoader 是 Pytorch（和 SenseParrots）负责数据读取的数据结构，在它的下层还有 Sampler 和 Dataset，他们的关系如下图所示：

图 1.Sampler 和 Dataset 的关系

Sampler 会根据用户的数据集进行采样，确定每个数据对应的 indices。Dataset 接收到 indices 之后会读取出图片，这时候的图片经过解码已经变成了数据表示，之后根据模型需要还会进行预处理，例如翻转、裁剪、归一化等等。经过预处理之后，数据才会作为模型的输入进入模型。

模型运行时的计算通常是在 GPU（或者其他类似的计算设备）进行的。这个时候，CPU 和磁盘 I/O 处在一个相对闲置的状态，而图像的读取和预处理最需要的就是 CPU 和 I/O。因此，我们自然而然就想到在这里设计并行方案：在模型拿到这一轮的数据进行计算时，我们同时去读取后面的轮次所需要的数据并进行预处理。这样的话，模型完成当前轮次的计算后，就能马上拿到数据开始下一轮次的计算而无需等待。于是，专门负责进行数据读取和预处理的 Worker 就这么出现了。

由于一个 Worker 可能还是无法完全用满 I/O，我们常常会根据 CPU 和 I/O 的负载情况，以及模型的实际情况来设置 Worker 的数目。Worker 进程（由于 Python 存在 GIL 锁，因此使用进程而不是线程）和主进程的协同如下图所示：

图 2.Worker 进程与主进程的协同

在上图中，主进程通过 Sampler 进行取样，发布读取任务到任务队列中。Worker 每完成一次读取，就从自己的任务队列中获取任务，根据任务的内容读取数据、进行预处理，并把数据放到数据队列中。主进程每次需要数据时从数据队列获取数据。这样，一个高效的数据读取器就完成了。

2 多进程的启动方式

Python 子进程有两种启动方式：fork 和 spawn。Spawn 会从头启动子进程再加载各种资源，启动慢但是更安全；fork 会拷贝当前内存空间，启动更快但是容易出现锁冲突之类的资源冲突。

于是，结合了 fork 和 spawn 优势的启动方式 forkserver 应运而生，原理如下图所示：

图 3. forkserver 的工作方式

正如上图，在某个进程需要频繁启动多个子进程时，forkerserver 可以先用 spawn 方法启动一个相对安全、纯净的子进程，再从这个子进程中 fork 进程，fork 出进程相当于是原进程的「孙子进程」。这种方法相当于是一次 spawn 加上多次 fork，能够兼具 spawn 的安全和 fork 的高效。

基于以上原因，我们在实际应用中，以 forkserver 为默认模式来启动 Worker，同时根据不同场景的稳定性和速度的要求，切换使用 spawn 和 fork 模式。

3 ShareMemory 的数据传输

进程间的通信方式有管道、命名管道、消息队列、共享内存……想必很多新手从业者都背过这样的八股。正如大家背过的那样，共享内存（ShareMemory）是最快的 IPC 方式。那么，我们自然就会想把它应用在 Dataloader 的加速上。

在 Worker 中，我们可以用共享内存来存放完成读取和预处理后的数据，然后把共享内存注册的文件描述符（fd）放到 Worker 和主进程的传输队列中供主进程获取。这样，队列中不传数据，只传 fd，就能变快很多。进一步的，我们还可以在子进程进行 ShareMemory 的缓存，减少申请 ShareMemory 所需的时间。

4 PinMemory 加速

PinMemory，即锁页内存，常常被用在 GPU 等计算硬件设备的加速上。平时我们使用内存时，系统会把内存数据交换到虚拟内存（磁盘）中，需要时再重新从磁盘加载，这样做可以让我们的内存上限比实际内存更大。而锁页内存就是系统在分配内存时锁定该页，让其不与磁盘交换。

我们读取的数据放到内存上之后，最终需要拷贝到 GPU（或者其他类似的计算设备）进行计算。英伟达等各个硬件厂商都有自己的 API 来注册锁页内存，这样能够实现内存的设备映射和传输并行，从而达到加速的效果。

5 Prefetch 预取

在很多时候，数据的大小规模、预处理的计算量、磁盘的 I/O 在模型多轮的运行中都是不稳定的。因此，每一次数据准备的时间也很可能是不稳定的。同样是 1s，我们有时候可能能取出 4 份数据，有时候只能取出 2 份。在这种背景下，我们设置数据的预取（也可以说是缓存），可以有效地避免数据读取的速度波动对整体速度的影响。

6 Worker 动态变化与最佳 Worker 数目搜索

在实际应用中，我们常常发现 Worker 数目并不是越多越好，也不是越少越好，而是有一个较优值。这是因为模型的训练计算部分往往也涉及各种翻译、调度等，不仅需要 GPU（或者其他类似的计算设备）资源，也需要 CPU 资源，过多的 Worker 数目反而会引发资源抢占，拖慢整体的训练速度。

因此，我们可以实现训练过程中 Worker 数目的动态变化，再配合合理的测速评估方案来形成工具，就能够自动搜索到最佳的 Worker 数目。当然，用户也可以经过一定的试验和调校来找到最合适的 Worker 数目。

7 PoolWorker

在一个训练任务中，我们往往会使用一个数据集进行多次的训练迭代（多 epoch 训练）。在这种情况下，如果每次 Dataloader 迭代完了数据就关闭 Worker，下一次新的循环迭代还得再启动 Worker，就显得非常浪费。

由此，我们可以通过 Worker 资源池来实现 Worker 资源的管理：如果资源池中有 Worker，那么想要启动新 Worker 时就可以直接从 Worker 资源池取出 Worker 使用；当数据读取任务全部完成时，我们会把 Worker 归还给资源池而不是关闭它。这样，我们就节省了 Worker 启动时的等待时间，提高了模型运行的运行速度。

8 DeviceTransform

当预处理相对复杂耗时的时候，也就是 CPU 的资源已经不足以让数据读取的一整个过程和模型网络计算部分完全并行重叠时，就会存在模型以及计算完成，GPU（或者其他类似的计算设备）在空等数据的情况。对此，我们可以把模型的预处理放到计算硬件设备上去来提高并行效率，避免计算资源浪费。概述中提到的英伟达推出的 DALI，就是把一些相对通用的预处理整合起来，提供了高性能的算子来方便用户在 GPU 上进行预处理。

9 其他

从加速的目的出发，数据读取还有许多功课可以做。例如，改用 lmdb 等其他数据格式来保存数据，使用 Memcached、Ceph 等进行存储上的优化，使用速度更快的图片解码库或者图片加载包等等……

在这里，我们只是对数据读取这一领域进行了一些简单的、总述性的介绍。实际上，数据读取还有很多有意思的方向值得挖掘，上文提到的各个方向也有很多细节实现值得探讨。让我们共同期待这个领域内有更多有意思的解决方案和工作成果涌现。