我想在我的输入管道中编写一个数据增强步骤,从概念上讲,我有两个数据集,它们可以作为一对输入到生成器中,在那里它们将产生一堆输出示例。
通过执行以下操作,我设法实现了这样的目标:
import tensorflow as tf
def gen(a, b):
for i in range(2):
yield str(a) + " " + str(b) + " " + str(i)
a = tf.data.Dataset.range(3)
b = tf.data.Dataset.range(3)
dataset = b.interleave(lambda x: a.interleave(lambda y: tf.data.Dataset.from_generator(gen,
output_types=(tf.string),
args=(x, y)),
num_parallel_calls = None))
for d in dataset:
print (d.numpy())
这会产生:
b'0 0 0'
b'0 1 0'
b'0 2 0'
b'0 0 1'
b'0 1 1'
b'0 2 1'
b'1 0 0'
b'1 1 0'
b'1 2 0'
b'1 0 1'
b'1 1 1'
b'1 2 1'
b'2 0 0'
b'2 1 0'
b'2 2 0'
b'2 0 1'
b'2 1 1'
b'2 2 1'
正如预期的那样。我的问题gen
是(在我的真实情况下)计算量很大的操作,所以我想尽可能使用并行调用。到目前为止,我尝试添加 num_parallel_calls 未能产生性能提升。
此外,如果它很重要,我的输入数据集来自一个TFRecordDataset
提供更多添加num_parallel_calls
选项的机会,即
raw_a = tf.data.TFRecordDataset(a_tfrecord_list)
a = raw_dataset.map(some_parsing_function)