python - 如何将多个 tfrecord 数据集合并为一个数据集？

Question

假设我有3个tfrecord文件，分别是neg.tfrecord, pos1.tfrecord, pos2.tfrecord。

我的batch size是500，包括300个neg数据，100个pos1数据，100个pos2数据。如何获得所需的 TFRecordDataset？

我将在 keras.fit()(急切执行)中使用此 TFRecordDataset 对象。

我的tensorflow版本是1.13.1。我在 tf.data.Dataset 中找到了 API ，比如interleave，concatenate，zip，但是好像不能解决我的问题。

之前尝试过获取每个数据集的iterator，拿到数据后再手动concat，但是效率很低，GPU利用率也不高。

在这个question ，我在下面使用 interleave:

tfrecord_files = ['neg.tfrecord', 'pos1.tfrecord', 'pos2.tfrecord']
dataset = tf.data.Dataset.from_tensor_slices(tfrecord_files)
def _parse(x):
    x = tf.data.TFRecordDataset(x)
    return x
dataset = dataset.interleave(_parse, cycle_length=4, block_length=1)
dataset = dataset.apply(tf.data.experimental.map_and_batch(_parse_image_function, 500))

我得到了这批:

neg pos1 pos2 neg pos1 pos2 ...............

但我想要的是:

neg neg neg pos1 pos2 neg neg neg pos1 pos2 .................

我该怎么办？

期待回答。

Answer 1

我用字符串数据复制了你所说的内容:

import tensorflow as tf

def string_data(s):
    return tf.sparse.to_dense(tf.strings.split([s]), default_value='')[0]

data = [' '.join(['neg'] * 30), ' '.join(['pos1'] * 10), ' '.join(['pos2'] * 10)]
step_sizes = tf.constant([3, 1, 1], dtype=tf.int64)
ds = (tf.data.Dataset.from_tensor_slices((data, step_sizes))
      .interleave(lambda d, s: (tf.data.Dataset.from_tensor_slices(string_data(d))
                                .batch(s)),
                  cycle_length=len(data))
      .flat_map(tf.data.Dataset.from_tensor_slices))
iter = ds.make_one_shot_iterator().get_next()

with tf.Session() as sess:
    while True:
        try:
            print(sess.run(iter).decode(), end=', ')
        except tf.errors.OutOfRangeError: break
    print()

输出:

neg, neg, neg, pos1, pos2, neg, neg, neg, pos1, pos2, neg, neg, neg, pos1, pos2, neg, neg, neg, pos1, pos2, neg, neg, neg, pos1, pos2, neg, neg, neg, pos1, pos2, neg, neg, neg, pos1, pos2, neg, neg, neg, pos1, pos2, neg, neg, neg, pos1, pos2, neg, neg, neg, pos1, pos2, 

在实际用例中，您会将 data 替换为文件名列表，将 tf.data.Dataset.from_tensor_slices(string_data(d)) 替换为 tf.data.TFRecordDataset(d)，但在其他方面应该类似。

编辑:我刚刚意识到你实际上想要一批以这种方式排序的所有元素，而不仅仅是一次一个元素，所以我想你必须在结束。