使用 tf.data 和叉积对两个数据集进行增强

Question

我想在我的输入管道中编写一个数据增强步骤，从概念上讲，我有两个数据集可以作为一对输入到生成器，在那里它们将产生一堆输出示例。

我已经通过以下方式实现了这样的目标：

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)

Answer 1

如果gen 是昂贵的部分，你应该尽量避免使用Dataset.from_generator。 Dataset.from_generator 需要为每个元素运行 Python 生成器，由于 Python 全局解释器锁，无法并行化。

您可以使用range + map 而不是from_generator 重写您的数据集：

import tensorflow as tf

a = tf.data.Dataset.range(3)
b = tf.data.Dataset.range(3)

def make_dataset(x, y):
  ds = tf.data.Dataset.range(2)
  ds = ds.map(lambda i: tf.strings.as_string(x) + " " + 
              tf.strings.as_string(y) + " " + tf.strings.as_string(i))
  return ds

dataset = b.interleave(lambda x: a.interleave(lambda y: make_dataset(x, y)))

for d in dataset:
    print (d.numpy())

要并行化昂贵的计算，请将num_parallel_calls=tf.data.experimental.AUTOTUNE 传递给map 和interleave 方法。