使用from_generator的Tensorflow数据管道生成数据的速度不够快

Question

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

Tensorflow 2.0，默认急切模式。

管道-读取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。因此，您必须将生成器重写为一次返回one