洗牌tfrecords文件

时间:2018-02-08 18:27:34

标签: tensorflow deep-learning tensorflow-datasets tfrecord

我有5个tfrecords文件,每个对象一个。虽然训练我想从所有5个tfrecords同样读取数据,即如果我的批量大小是50,我应该得到1 tfrecord文件10个样本,从第二tfrecord文件10个样品等。目前,它只是从所有三个文件中顺序读取,即我从同一记录中获得50个样本。有没有办法从不同的tfrecords文件中采样?

1 个答案:

答案 0 :(得分:7)

我建议您在tf.data上阅读@mrry的tutorial。在幻灯片42中,他解释了如何使用tf.data.Dataset.interleave()同时读取多个tfrecord文件。

例如,如果您有5个文件,包含:

file0.tfrecord: [0, 1]
file1.tfrecord: [2, 3]
file2.tfrecord: [4, 5]
file3.tfrecord: [6, 7]
file4.tfrecord: [8, 9]

您可以像这样编写数据集:

files = ["file{}.tfrecord".format(i) for i in range(5)]
files = tf.data.Dataset.from_tensor_slices(files)
dataset = files.interleave(lambda x: tf.data.TFRecordDataset(x),
                           cycle_length=5, block_length=1)

dataset = dataset.map(_parse_function)  # parse the record

interleave的参数是: - cycle_length:要同时读取的文件数。如果要从所有文件中读取以创建批处理,请将其设置为文件数(在您的情况下,这是您应该执行的操作,因为每个文件包含一种类型的标签) - block_length:每次我们从文件中读取时,都会从此文件中读取block_length个元素

我们可以测试它是否按预期工作:

iterator = dataset.make_one_shot_iterator()
x = iterator.get_next()

with tf.Session() as sess:
    for _ in range(num_samples):
        print(sess.run(x))

将打印:

0
2
4
6
8
1
3
5
7
9