TFRecordReader 似乎非常慢，并且多线程读取不起作用

Question

我的训练过程对 train&eval 数据集使用 tfrecord 格式。

我测试了 reader 的 benchmark，只有 8000 条记录/秒。和 io 速度（见 iotop 命令）只有 400KB-500KB/s。

我这里用的是cpp版本的protobuf

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/g3doc/get_started/os_setup.md#protobuf-library-related-issues

如果可能，请提供一个可重现的最小示例（我们通常没有时间阅读您的数百行代码）

def read_and_decode(filename_queue):
     reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)
    return serialized_example
  serialized_example = read_and_decode(filename_queue)
  batch_serialized_example = tf.train.shuffle_batch(
      [serialized_example],
      batch_size=batch_size,
      num_threads=thread_number,
      capacity=capacity,
      min_after_dequeue=min_after_dequeue)
  features = tf.parse_example(
      batch_serialized_example,
      features={
          "label": tf.FixedLenFeature([], tf.float32),
          "ids": tf.VarLenFeature(tf.int64),
          "values": tf.VarLenFeature(tf.float32),
      })

您还尝试过哪些其他尝试性解决方案？

我尝试在 tf.train.shuffle_batch 中设置 num_threads 但不起作用。

似乎设置为2线程时，它以8000records / s的速度工作，当增加线程数时，它变得更慢。 （我删除了所有消耗 cpu 的操作。只需读取数据。）

我的服务器是 24 核 CPU。

Answer 1

这里的问题是每个session.run 都有固定的开销开销，并且向队列填充许多微小的示例会很慢。

特别是，每个 session.run 大约 100-200 微秒，因此您每秒只能进行大约 5k-10k session.run 调用。

如果进行 Python 分析 (python -m cProfile)，这个问题很明显，但很难看出是从时间线配置文件还是 CPU 配置文件开始。

解决方法是使用enqueue_many 将内容批量添加到您的队列中。我从 https://gist.github.com/ericyue/7705407a88e643f7ab380c6658f641e8 中获取了您的基准，并对其进行了修改，以便在每次 .run 调用时将许多项目排入队列，这样可以加快 10 倍的速度。

修改是修改tf.batch调用如下：

if enqueue_many:
    reader = tf.TFRecordReader(options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB))
    queue_batch = []
    for i in range(enqueue_many_size):
        _, serialized_example = reader.read(filename_queue)
        queue_batch.append(serialized_example)
    batch_serialized_example = tf.train.shuffle_batch(
        [queue_batch],
        batch_size=batch_size,
        num_threads=thread_number,
        capacity=capacity,
        min_after_dequeue=min_after_dequeue,
        enqueue_many=True)

如需完整源代码，请查看此处： https://github.com/yaroslavvb/stuff/blob/master/ericyue-slowreader/benchmark.py

很难对其进行优化以使其运行得更快，因为现在大部分时间都花在了队列操作上。查看 stripped down 版本，它只是将整数添加到队列中，您也可以获得类似的速度，并且查看时间线，时间花在出队操作上。

每个出队操作大约需要 60 微秒，但平均有 5 个并行运行，因此每个出队可获得 12 微秒。这意味着在最好的情况下，您每秒将获得

Answer 2

这是基于雅罗斯拉夫答案的简单加速构建：

Tensorflow 有一个内置函数tf.TFRecordReader.read_up_to，它会在每个session.run() 调用中读取多条记录，从而消除多次调用造成的额外开销。

enqueue_many_size = SOME_ENQUEUE_MANY_SIZE
reader = tf.TFRecordReader(options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB))
_, queue_batch = reader.read_up_to(filename_queue, enqueue_many_size)
batch_serialized_example = tf.train.shuffle_batch(
    [queue_batch],
    batch_size=batch_size,
    num_threads=thread_number,
    capacity=capacity,
    min_after_dequeue=min_after_dequeue,
    enqueue_many=True)

与 Yaroslav 的回答一样，您需要设置 enqueue_many=True 以便批处理函数知道它正在接受多条记录。

这在我的用例中非常快。

Answer 3

雅罗斯拉夫回答的附录： 您可以使用tf.python_io.tf_record_iterator 遍历示例，以便将它们附加到您可以使用enqueue_many=true 传递给tf.train.shuffle_batch 的列表中：

queue_batch = []
for serialized_example in tf.python_io.tf_record_iterator(filename,options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)):
    queue_batch.append(serialized_example)
batch_serialized_example = tf.train.shuffle_batch(
    [queue_batch],
    batch_size=batch_size,
    num_threads=thread_number,
    capacity=capacity,
    min_after_dequeue=min_after_dequeue,
    enqueue_many=True)

似乎尝试使用reader.read() 迭代示例将导致每批读取一次。即第 n 批将是第 n 条记录的 batch_num 副本，而不是 batch_num 许多唯一记录。