Python 将.tfrecords文件拆分为多个.tfrecords文件

有没有办法直接将.tfrecords文件拆分为多个.tfrecords文件，而不回写每个数据集示例？

您可以使用如下函数：

import tensorflow as tf

def split_tfrecord(tfrecord_path, split_size):
    with tf.Graph().as_default(), tf.Session() as sess:
        ds = tf.data.TFRecordDataset(tfrecord_path).batch(split_size)
        batch = ds.make_one_shot_iterator().get_next()
        part_num = 0
        while True:
            try:
                records = sess.run(batch)
                part_path = tfrecord_path + '.{:03d}'.format(part_num)
                with tf.python_io.TFRecordWriter(part_path) as writer:
                    for record in records:
                        writer.write(record)
                part_num += 1
            except tf.errors.OutOfRangeError: break

例如，要将文件

my_records.tfrecord

拆分为各100条记录的部分，您可以执行以下操作：

split_tfrecord(my_records.tfrecord, 100)

---

这将创建多个较小的记录文件

my_records.tfrecord.000

，

my_records.tfrecord.001

，等等。在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))

使用

.batch（）

而不是

.shard（）

避免在数据集上多次迭代 与使用

tf.data.Dataset.shard（）

）相比，一种性能更高的方法是使用批处理：

将tensorflow导入为tf
ITEMS_PER_FILE=100#假设我们每个.tfrecord文件保存100个项目
原始数据集=tf.data.TFRecordDataset（'in.tfrecord'）
批次_idx=0
对于原始数据集中的批处理。批处理（每个文件中的项目）：
#将“batch”转换回“Dataset”，假设batch是“张量”的“元组”`
batch\u ds=tf.data.Dataset.from\u tensor\u切片（元组（[*batch]））
filename=f'out.tfrecord.{batch_idx:03d}'
writer=tf.data.experimental.TFRecordWriter（文件名）
writer.write（批处理）
批次_idx+=1

TensorFlow 2.x非常有效的方法 正如@yongjieyongjie所提到的，您应该使用

.batch（）

而不是

.shard（）

，以避免根据需要在数据集上进行更频繁的迭代。 但如果您有一个非常大的数据集，内存太大，它将失败（但没有错误），只会给您一些文件和原始数据集的一小部分

首先，您应该对数据集进行批处理，并使用每个文件想要的记录量作为批处理大小（我假设您的数据集已采用序列化格式，否则请参见）

dataset=dataset.batch（每个文件的项目）

接下来要做的是使用生成器以避免内存不足

def write_生成器（）：
i=0
迭代器=iter（数据集）
optional=迭代器。将_next_作为_optional（）获取
而可选。具有_值（）.numpy（）：
ds=可选。获取_值（）
optional=迭代器。将_next_作为_optional（）获取
batch\u ds=tf.data.Dataset.from\u tensor\u切片（ds）
writer=tf.data.experimental.TFRecordWriter（保存到+“\\”+name+“-”+str（i）+“.tfrecord”，压缩类型='GZIP'）压缩类型='GZIP'
i+=1
产量批次，作者，i
返回

现在只需在正常的for循环中使用生成器

对于数据，wri，i在write_生成器（）中：
开始时间=time.time（）
写入（数据）
打印（“所需时间：”，Time.Time（）

只要一个文件适合内存中的原始文件，就可以正常工作。

不均匀分割 如果要将文件平均分割为大小相同的文件，则大多数其他答案都有效。这将适用于不均匀分割：

# `splits` is a list of the number of records you want in each output file
def split_files(filename: str, splits: List[int]) -> None:
    dataset: tf.data.Dataset = tf.data.TFRecordDataset(filename)
    rec_counter: int = 0

    # An extra iteration over the data to get the size
    total_records: int = len([r for r in dataset])
    print(f"Found {total_records} records in source file.")

    if sum(splits) != total_records:
        raise ValueError(f"Sum of splits {sum(splits)} does not equal "
                         f"total number of records {total_records}")

    rec_iter:Iterator = iter(dataset)
    split: int
    for split_idx, split in enumerate(splits):
        outfile: str = filename + f".{split_idx}-{split}"
        with tf.io.TFRecordWriter(outfile) as writer:
            for out_idx in range(split):
                rec: tf.Tensor = next(rec_iter, None)
                rec_counter +=1
                writer.write(rec.numpy())
        print(f"Finished writing {split} records to file {split_idx}")

虽然从技术上讲，我认为OP询问

时没有回写每个数据集示例（这就是它所做的），但这至少是在没有反序列化每个示例的情况下进行的

对于非常大的文件来说有点慢。可能有一种方法可以修改其他一些基于批处理的答案，以便使用批处理输入读取，但仍然写入不均匀的拆分，但我没有尝试。
在N个拆分中拆分（在tensorflow 1.13.1中测试）

import os
import hashlib
import tensorflow as tf
from tqdm import tqdm


def split_tfrecord(tfrecord_path, n_splits):
    dataset = tf.data.TFRecordDataset(tfrecord_path)
    outfiles=[]
    for n_split in range(n_splits):
        output_tfrecord_dir = f"{os.path.splitext(tfrecord_path)[0]}"
        if not os.path.exists(output_tfrecord_dir):
            os.makedirs(output_tfrecord_dir)
        output_tfrecord_path=os.path.join(output_tfrecord_dir, f"{n_split:03d}.tfrecord")
        out_f = tf.io.TFRecordWriter(output_tfrecord_path)
        outfiles.append(out_f)

    for record in tqdm(dataset):
        sample = tf.train.Example()
        record = record.numpy()
        sample.ParseFromString(record)

        idx = int(hashlib.sha1(record).hexdigest(),16) % n_splits
        outfiles[idx].write(example.SerializeToString())

    for file in outfiles:
        file.close()