1

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

通过执行以下操作,我设法实现了这样的目标:

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)
4

1 回答 1

0

如果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.AUTOTUNEmapandinterleave方法。

于 2019-12-27T18:31:59.437 回答