python - Split .tfrecords file into many .tfrecords files

Question

Is there any way to split .tfrecords file into many .tfrecords files directly, without writing back each Dataset example ?

Answer 1

In tensorflow 2.0.0, this will work:

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))