如何将多个 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)，但在其他方面应该类似。

编辑：我刚刚意识到你实际上想要一批以这种方式排序的所有元素，而不仅仅是一次一个元素，所以我想你必须在最后添加另一个 batch 调用.

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

How to combine multiple tfrecord datasets into one dataset?

python

deep-learning

tensorflow

tfrecord