TensorFlow Data饥饿的GPU

Question

我需要帮助优化自定义TensorFlow模型。我有一个40GB的ZLIB压缩.TFRecords文件，包含我的训练数据。每个样本由两个384x512x3图像和一个384x512x2矢量场组成。我正在加载我的数据如下：

    num_threads = 16
    reader_kwargs = {'options': tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)}
    data_provider = slim.dataset_data_provider.DatasetDataProvider(
                        dataset,
                        num_readers=num_threads,
                        reader_kwargs=reader_kwargs)
    image_a, image_b, flow = data_provider.get(['image_a', 'image_b', 'flow'])

    image_as, image_bs, flows = tf.train.batch(
        [image_a, image_b, flow],
        batch_size=dataset_config['BATCH_SIZE'], # 8
        capacity=dataset_config['BATCH_SIZE'] * 10,
        num_threads=num_threads,
        allow_smaller_final_batch=False)

然而，我只得到大约0.25到0.30全球步数/秒。 （SLOW！）

这是我用于并行阅读器的TensorBoard短划线。它始终保持在99％-100％。

我随着时间的推移绘制了GPU使用率（％/秒）。它看起来数据饥饿，但我不知道如何解决这个问题。我尝试过增加/减少线程的数量，但它似乎没有什么区别。我正在使用具有4个CPU和61GB RAM的NVIDIA K80 GPU进行培训。

如何让这列火车更快？

Answer 1

如果您的示例很小，那么使用DataSetProvider将无法获得令人满意的结果。它一次只读一个例子，可以是一个瓶颈。我已经添加了feature request on github。

与此同时，您必须使用自己的使用read_up_to的输入队列进行滚动：

  batch_size = 10000
  num_tfrecords_at_once = 1024
  reader = tf.TFRecordReader()
  # Here's where the magic happens:
  _, records = reader.read_up_to(filename_queue, num_tfrecords_at_once)

  # Batch records with 'enqueue_many=True'
  batch_serialized_example = tf.train.shuffle_batch(
      [records],
      num_threads=num_threads,
      batch_size=batch_size,
      capacity=10 * batch_size,
      min_after_dequeue=2 * batch_size,
      enqueue_many=True)

  parsed = tf.parse_example(
      batch_serialized_example,
      features=whatever_features_you_have)
  # Use parsed['feature_name'] etc. below