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

时间:2017-01-14 07:33:20

标签: python tensorflow

我的训练过程使用tfrecord格式进行训练和评估数据集。

我测试了阅读器的基准,只有8000记录/秒。和io速度(参见iotop命令)只需400KB-500KB / s。

我在这里使用protobuf的cpp版本

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,当放大线程数时,它会变慢。 (我删除所有花费cpus的操作。只需读取数据。)

我的服务器是24核心cpu。

3 个答案:

答案 0 :(得分:8)

这里的问题是每个session.run都有一个固定的成本开销,并且用队列中的许多小例子填充队列会很慢。

特别是,每个session.run大约是100-200 usec,因此每秒只能进行大约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版本,您也可以获得类似的速度,并且查看时间轴,将时间用于出列操作。

enter image description here

每个dequeue op大约需要60个usec,但平均有5个并行运行,所以每个dequeue得到12个usec。这意味着在最好的情况下你将获得每秒200万的例子。

答案 1 :(得分:5)

这是雅罗斯拉夫答案的简单加速建设:

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,以便批处理函数知道它正在接受多个记录。

我的用例非常快。

答案 2 :(得分:1)

雅罗斯拉夫答复的附录: 您可以使用tf.python_io.tf_record_iterator来迭代这些示例,然后将它们附加到一个列表,您可以使用tf.train.shuffle_batch将其传递给enqueue_many=true

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个许多唯一记录。