使用Ray創建高效的深度學習資料管道

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

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))

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 <p>现在，如果测量等待下一个小批次数据所花费的时间，得到的平均时间为0.08秒。速度提高了近7倍，但理想情况下，希望这个时间接近0。</p><p>如果进行分析，可以发现相当多的时间都花在了准备数据的反序列化上。在多进程生成器中，生产者进程会返回大型NumPy数组，这些数组需要进行准备，然后在主进程中进行反序列化。那么在进程间传递大型数组时，如何提高效率呢？</p><h2>4、使用Ray并行化数据管道</h2><p>这就是Ray发挥作用的地方。Ray是一个用于在Python中运行分布式计算的框架。它带有一个共享内存对象存储区，可在不同进程间高效地传输对象。特别的是，在不进行任何序列化和反序列化的情况下，对象存储区中的Numpy数组可在同一节点上的worker之间共享。Ray还可以轻松实现数据加载在多台机器上的扩展，并使用Apache Arrow高效地序列化和反序列化大型数组。</p><p>Ray带有一个实用函数from_iterators，可以创建并行迭代器，开发者可以用它包装data_generator生成器函数。</p><pre class="brush:php;toolbar:false">import raydef ray_generator():num_parallel_processes = cpu_count()return ray.util.iter.from_iterators([data_generator]*num_parallel_processes).gather_async()