带有 from_generator 的 TensorFlow 数据管道生成数据的速度不够快

Question

我正在使用 tf.data.Dataset 构建训练管道，但我似乎无法获得速度。我在管道中使用 from_generator （见下文），我怀疑 python 全局锁是问题的一部分。我正在努力找出一种可以加快速度的设计。我在交错中尝试了 num_parallel_calls、cycle_length、block_length 的不同变体，但它们似乎都没有帮助。欣赏任何见解。

Tensorflow 2.0，默认 Eager 模式。

管道 - 读取 TFRecords 并将数据提供给训练

我有两个生成器函数，我覆盖它们以分散工作负载。

• generator_step1 - 加载 tf_record 文件，使用 map 解析数据
• generator_step2 - 处理数据并产生批量数据

管道伪代码

dataset = tf.data.Dataset.from_tensor_slices(tf_records_files_list)

dataset = dataset.interleave(lambda x: tf.data.Dataset.from_generator(generator_step1, 
                                output_types=(tf.string, tf.float32), args=(x,batch_size,)),
                                cycle_length=2,
                                block_length=1)

dataset = dataset.interleave(lambda arg1, arg2: tf.data.Dataset.from_generator(generator_step2, 
                                output_types=(tf.float32, tf.int16), args=(arg1, arg2,)),
                                cycle_length=2,
                                block_length=1)

dataset = dataset.prefetch(5)

这在一台功能强大的机器上运行，我认为 I/O 或 CPU 不是瓶颈。关于如何进一步改进或优化这一点的任何想法？欣赏任何见解。

Answer 1

好的，我认为你的设计不是最优的。首先创建一个tf.dataset，然后为每个批次创建另一个tf.dataset，这会产生大量开销。而不是interleave，我建议使用.map。因此，您必须将生成器重写为一次返回 一个 样本的函数。之后你应该通过tf.data 进行批处理。此外，您应该在映射函数中使用 tensorflow 函数（例如 tf.image.decode_iamge）。

dataset = tf.data.Dataset.from_tensor_slices(tf_records_files_list)

dataset = dataset.map(load_func, num_parralel_calls=24)   

dataset = dataset.map(preprocess_func, num_parralel_calls=24)

dataset = dataset.batch(batch_size).prefetch(5)

这应该可以完成工作。