python - 将 .tfrecords 文件拆分为多个 .tfrecords 文件

Question

有什么方法可以直接将 .tfrecords 文件拆分为多个 .tfrecords 文件，而无需写回每个数据集示例？

Answer 1

在 tensorflow 2.0.0 中，这将起作用:

import tensorflow as tf

raw_dataset = tf.data.TFRecordDataset("input_file.tfrecord")

shards = 10

for i in range(shards):
    writer = tf.data.experimental.TFRecordWriter(f"output_file-part-{i}.tfrecord")
    writer.write(raw_dataset.shard(shards, i))