TensorFlow 数据匮乏的 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 步/秒的全局步数。 （慢！）

这是我用于并行阅读器的 TensorBoard dash。它始终保持在 99%-100%。

我绘制了一段时间内的 GPU 使用情况（每秒百分比）。它看起来数据匮乏，但我不知道如何解决这个问题。我试过增加/减少线程数，但似乎没有什么区别。我正在使用具有 4 个 CPU 和 61GB 内存的 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