tensorflow - 并行化 tf.data.Dataset.from_generator

Question

我有一个不平凡的输入管道， from_generator 非常适合...

dataset = tf.data.Dataset.from_generator(complex_img_label_generator,
                                        (tf.int32, tf.string))
dataset = dataset.batch(64)
iter = dataset.make_one_shot_iterator()
imgs, labels = iter.get_next()

其中complex_img_label_generator动态生成图像并返回一个表示(H, W, 3)图像和一个简单的string标签的numpy数组。我无法将处理表示为从文件读取和 tf.image 操作。

我的问题是如何并行化生成器？我如何让 N 个这些生成器在它们自己的线程中运行。

一种想法是使用 dataset.map 和 num_parallel_calls 来处理线程；但 map 在张量上运行...另一个想法是创建多个生成器，每个生成器都有自己的预取并以某种方式连接它们，但我不知道如何连接N个生成器流？

我可以遵循任何典型的例子吗？

Answer 1

事实证明，如果我使生成器超轻量级(仅生成元数据)，然后将实际的重型照明移动到无状态函数中，我可以使用 Dataset.map 。这样我就可以使用 py_func 将繁重的工作与 .map 并行化。

作品；但感觉有点笨拙...如果能够将 num_parallel_calls 添加到 from_generator 就太好了:)

def pure_numpy_and_pil_complex_calculation(metadata, label):
  # some complex pil and numpy work nothing to do with tf
  ...

dataset = tf.data.Dataset.from_generator(lightweight_generator,
                                         output_types=(tf.string,   # metadata
                                                       tf.string))  # label

def wrapped_complex_calulation(metadata, label):
  return tf.py_func(func = pure_numpy_and_pil_complex_calculation,
                    inp = (metadata, label),
                    Tout = (tf.uint8,    # (H,W,3) img
                            tf.string))  # label
dataset = dataset.map(wrapped_complex_calulation,
                      num_parallel_calls=8)

dataset = dataset.batch(64)
iter = dataset.make_one_shot_iterator()
imgs, labels = iter.get_next()