tensorflow - TFRecord IO 比 python hdf5 阅读器慢，如何提高它的速度？

Question

我按照官方 TF 指南使用tf.data.DatasetAPI 来构建数据管道，但我发现它比我使用 hdf5 的 python 数据管道慢 2 倍。

这是我的实验结果：

TFRecordDataset：66万样本/秒
hdf5 reader + feed_dict 到占位符：110 万样本/秒

我的实验设置是：
- 批量大小 = 1000，
- 每 1000 万个样本打印一次日志以显示时间，
- 运行一个只将数据作为输入并且不计算任何东西的假模型。

更糟糕的是，当我设置batch size=10000时，它变成了：

TFRecordDataset：76万样本/秒
hdf5 reader + feed_dict 到占位符：210 万样本/秒

这是我阅读 tfrecord 的代码

def parse_tfrecord_batch(record_batch):
    FEATURE_LEN = 29
    dics = {
        'label': tf.FixedLenFeature(shape=(), dtype=tf.int64),
        'feature_id': tf.FixedLenFeature(shape=(FEATURE_LEN,), dtype=tf.int64),
        'feature_val': tf.FixedLenFeature(shape=(FEATURE_LEN,), dtype=tf.float32)
    }
    parsed_example = tf.parse_example(record_batch, dics)
    return parsed_example['feature_id'], parsed_example['feature_val'], parsed_example['label']

class fakeModel:
    def __init__(self, train_filenames):
        self.graph = tf.Graph()
        with self.graph.as_default():
            dataset = tf.data.TFRecordDataset(train_filenames)
            dataset = dataset.repeat()
            dataset = dataset.batch(1000)
            dataset = dataset.map(parse_tfrecord_batch, num_parallel_calls=1)
            dataset = dataset.prefetch(1000)
            self.iterator = dataset.make_initializable_iterator()
            self.id, self.wt, self.label = self.iterator.get_next()

            self.train_preds = tf.identity(self.lbl_hldr)

我已将其num_parallel_calls调至 2、10。无法正常工作。

我还将prefetch(n)从 1 调整到 1000，这几乎没有改进。

我的问题是：
有什么方法可以改善我的 tfrecord 数据管道？我的代码中是否缺少某些内容？

感谢它的任何帮助。

标签： tensorflowhdf5tensorflow-datasetstfrecord