Tensorflow：在CPU上的多个线程中加载数据

我有一个 python 课程SceneGenerator它有多个用于预处理的成员函数和一个生成器函数generate_data()。基本结构是这样的：

class SceneGenerator(object):
    def __init__(self):
       # some inits

    def generate_data(self):
        """
        Generator. Yield data X and labels y after some preprocessing
        """
        while True:
            # opening files, selecting data
            X,y = self.preprocess(some_params, filenames, ...)            

            yield X, y

我使用keras model.fit_generator()函数中的类成员函数sceneGenerator.generate_data()从磁盘读取数据，对其进行预处理并产生它。在 keras 中，这是在多个 CPU 线程上完成的，如果workers的参数model.fit_generator()设置为 > 1。

我现在想用同样的SceneGenerator张量流中的类。我目前的做法是这样的：

sceneGenerator = SceneGenerator(some_params)
for X, y in sceneGenerator.generate_data():

    feed_dict = {ops['data']: X,
                 ops['labels']: y,
                 ops['is_training_pl']: True
                 }
    summary, step, _, loss, prediction = sess.run([optimization_op, loss_op, pred_op],
                                                  feed_dict=feed_dict)

然而，这很慢并且不使用多线程。我找到了tf.data.Dataset https://www.tensorflow.org/versions/master/api_docs/python/tf/data/Datasetapi 与一些文档 https://www.tensorflow.org/versions/master/programmers_guide/datasets，但我未能实现这些方法。

Edit:请注意，我不处理图像，因此带有文件路径等的图像加载机制在这里不起作用。 我的SceneGenerator从 hdf5 文件加载数据。但不是完整的数据集，而是 - 根据初始化参数 - 仅数据集的一部分。我很想保持生成器功能不变，并了解如何将该生成器直接用作张量流的输入并在 CPU 上的多个线程上运行。将 hdf5 文件中的数据重写为 csv 并不是一个好的选择，因为它会重复大量数据。

Edit 2:：我认为类似的东西可能会有所帮助：并行化 tf.data.Dataset.from_generator https://stackoverflow.com/questions/47086599/parallelising-tf-data-dataset-from-generator

假设您使用的是最新的 Tensorflow（撰写本文时为 1.4），您可以保留生成器并使用tf.data.* https://www.tensorflow.org/api_docs/python/tf/dataAPI如下（我为线程数、预取缓冲区大小、批量大小和输出数据类型选择任意值）：

NUM_THREADS = 5
sceneGen = SceneGenerator()
dataset = tf.data.Dataset.from_generator(sceneGen.generate_data, output_types=(tf.float32, tf.int32))
dataset = dataset.map(lambda x,y : (x,y), num_parallel_calls=NUM_THREADS).prefetch(buffer_size=1000)
dataset = dataset.batch(42)
X, y = dataset.make_one_shot_iterator().get_next()

为了表明它实际上是从生成器中提取的多个线程，我修改了您的类，如下所示：

import threading    
class SceneGenerator(object):
  def __init__(self):
    # some inits
    pass

  def generate_data(self):
    """
    Generator. Yield data X and labels y after some preprocessing
    """
    while True:
      # opening files, selecting data
      X,y = threading.get_ident(), 2 #self.preprocess(some_params, filenames, ...)            
      yield X, y

这样，创建一个 Tensorflow 会话并获取一批即可显示获取数据的线程的线程 ID。在我的电脑上，运行：

sess = tf.Session()
print(sess.run([X, y]))

prints

[array([  8460.,   8460.,   8460.,  15912.,  16200.,  16200.,   8460.,
         15912.,  16200.,   8460.,  15912.,  16200.,  16200.,   8460.,
         15912.,  15912.,   8460.,   8460.,   6552.,  15912.,  15912.,
          8460.,   8460.,  15912.,   9956.,  16200.,   9956.,  16200.,
         15912.,  15912.,   9956.,  16200.,  15912.,  16200.,  16200.,
         16200.,   6552.,  16200.,  16200.,   9956.,   6552.,   6552.], dtype=float32),
 array([2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
        2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2])]

Note：您可能想尝试删除map调用（我们仅用于多线程）并检查是否prefetch的缓冲区足以消除输入管道中的瓶颈（即使只有一个线程，输入预处理通常比实际图形执行速度更快，因此缓冲区足以使预处理尽可能快地进行）。