使用Ray创建高效的深度学习数据管道

深度学习模型训练所需的gpu功能强大，但价格昂贵。为了充分利用gpu，开发人员需要一个高效的数据传输通道，以便在gpu准备好计算下一个训练步骤时，能够迅速将数据传输到gpu。使用ray能够显著提高数据传输通道的效率

1、训练数据管道的结构

首先，我们来看一下模型训练的伪代码

for step in range(num_steps):sample, target = next(dataset) # 步骤1train_step(sample, target) # 步骤2

在步骤1中，获取下一个小批量的样本和标签。在步骤2中，它们被传递给train_step函数，该函数会将它们复制到GPU上，执行前向传递和反向传递以计算损失和梯度，并更新优化器的权重。

请详细了解步骤1。当数据集太大无法放入内存时，步骤1将从磁盘或网络中获取下一个小批量数据。此外，步骤1还包括一定量的预处理。输入数据必须被转换为数字张量或张量集合，然后再馈送给模型。在某些情况下，还会在传递给模型之前对张量进行其他转换，如归一化、绕轴旋转、随机打乱等

如果工作流程是严格按顺序执行的，即先执行步骤1，然后再执行步骤2，那么模型将始终需要等待下一批数据的输入、输出和预处理操作。GPU将无法得到有效利用，它将在加载下一个小批量数据时处于空闲状态。

为了解决这个问题，可以将数据管道视为生产者——消费者的问题。数据管道生成小批量数据并写入有界缓冲区。模型/GPU从缓冲区中消费小批量数据，执行前向/反向计算并更新模型权重。如果数据管道能够以模型/GPU消费的速度快速生成小批量数据，那么训练过程将会非常高效。

☞☞☞AI 智能聊天, 问答助手, AI 智能搜索, 免费无限量使用 DeepSeek R1 模型☜☜☜

图片

2、Tensorflow tf.data API

Tensorflow tf.data API提供了一组丰富的功能，可用于高效创建数据管道，使用后台线程获取小批量数据，使模型无需等待。仅仅预先获取数据还不够，如果生成小批量数据的速度比GPU消费数据的速度慢，那么就需要使用并行化来加快数据的读取和转换。为此，Tensorflow提供了交错功能以利用多个线程并行读取数据，以及并行映射功能使用多个线程对小批量数据进行转换。

由于这些API是基于多线程的，所以可能会受到Python全局解释器锁（GIL）的限制。Python的GIL限制了一次只能运行单个线程的字节码。如果在管道中使用纯TensorFlow代码，通常不会受到这种限制，因为TensorFlow核心执行引擎在GIL的范围之外工作。但是，如果使用的第三方库没有解除GIL限制或者使用Python进行大量计算，那么依赖多线程来并行化管道就不可行

3、使用多进程并行化数据管道

考虑以下生成器函数，该函数模拟加载和执行一些计算以生成小批量数据样本和标签。

def data_generator():for _ in range(10):# 模拟获取# 从磁盘/网络time.sleep(0.5)# 模拟计算for _ in range(10000):passyield (np.random.random((4, 1000000, 3)).astype(np.float32), np.random.random((4, 1)).astype(np.float32))

接下来，在虚拟的训练管道中使用该生成器，并测量生成小批量数据所花费的平均时间。

generator_dataset = tf.data.Dataset.from_generator(data_generator,output_types=(tf.float64, tf.float64),output_shapes=((4, 1000000, 3), (4, 1))).prefetch(tf.data.experimental.AUTOTUNE)st = time.perf_counter()times = []for _ in generator_dataset:en = time.perf_counter()times.append(en - st)# 模拟训练步骤time.sleep(0.1)st = time.perf_counter()print(np.mean(times))

据观察，平均耗时约为0.57秒（在配备Intel Core i7处理器的Mac笔记本电脑上测量）。如果这是一个真实的训练循环，GPU的利用率将相当低，它只需花费0.1秒进行计算，然后闲置0.57秒等待下一个批次数据。

为了加快数据加载速度，可以使用多进程生成器。

from multiprocessing import Queue, cpu_count, Processdef mp_data_generator():def producer(q):for _ in range(10):# 模拟获取# 从磁盘/网络time.sleep(0.5)# 模拟计算for _ in range(10000000):passq.put((np.random.random((4, 1000000, 3)).astype(np.float32),np.random.random((4, 1)).astype(np.float32)))q.put("DONE")queue = Queue(cpu_count()*2)num_parallel_processes = cpu_count()producers = []for _ in range(num_parallel_processes):p = Process(target=producer, args=(queue,))p.start()producers.append(p)done_counts = 0while done_counts < num_parallel_processes:msg = queue.get()if msg == "DONE":done_counts += 1else:yield msgqueue.join()

现在，如果测量等待下一个小批次数据所花费的时间，得到的平均时间为0.08秒。速度提高了近7倍，但理想情况下，希望这个时间接近0。

如果进行分析，可以发现相当多的时间都花在了准备数据的反序列化上。在多进程生成器中，生产者进程会返回大型NumPy数组，这些数组需要进行准备，然后在主进程中进行反序列化。那么在进程间传递大型数组时，如何提高效率呢？

4、使用Ray并行化数据管道

这就是Ray发挥作用的地方。Ray是一个用于在Python中运行分布式计算的框架。它带有一个共享内存对象存储区，可在不同进程间高效地传输对象。特别的是，在不进行任何序列化和反序列化的情况下，对象存储区中的Numpy数组可在同一节点上的worker之间共享。Ray还可以轻松实现数据加载在多台机器上的扩展，并使用Apache Arrow高效地序列化和反序列化大型数组。

Ray带有一个实用函数from_iterators，可以创建并行迭代器，开发者可以用它包装data_generator生成器函数。

import raydef ray_generator():num_parallel_processes = cpu_count()return ray.util.iter.from_iterators([data_generator]*num_parallel_processes).gather_async()

使用ray_generator，测量等待下一个小批量数据所花费的时间为0.02秒，比使用多进程处理的速度提高了4倍。

# tensorflow # 区中 # 时间为 # 然后再 # 情况下 # 可在 # 多个 # 加载 # 序列化 # 小批量 # 深度学习 # apache # 对象 # 多线程 # 线程 # 循环 # numpy # 分布式 # Python