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方法。

将tf.data与叉积扩充为两个数据集

1 个答案: