Tensorflow - 使用 parallel_interleave 从多个 tfrecords 中读取不同的 block_lengths?

Tensorflow - Read different block_lengths from multiple tfrecords with parallel_interleave?

我正在尝试读取三个不同长度的大 tfrecords,并像这样并行读取它们:

files = [ filename1, filename2, filename3 ]

data = tf.data.TFRecordDataset(files)

data = data.apply(
    tf.contrib.data.parallel_interleave(
        lambda filename: tf.data.TFRecordDataset(data),
        cycle_length=3,block_length = [10,5,3]))

data = data.shuffle(
    buffer_size = 100)

data = data.apply(
    tf.contrib.data.map_and_batch(
        map_func=parse, 
        batch_size=100))

data = data.prefetch(10)

,但 TensorFlow 不允许每个文件源具有不同的块长度:

InvalidArgumentError: block_length must be a scalar

我可以创建三个具有不同小批量大小的不同数据集,但这需要 3 倍的资源,而且这不是我的机器限制所提供的选项。

可能的解决方案是什么?

这是答案,我想出了在我的限制范围内如何做到这一点。

为每个文件制作数据集,为每个文件定义每个小批量大小,并将 get_next() 输出连接在一起。这适合我的机器并且运行高效。