如何批量写入tf.train.Example()到TFRecords以降低大数据处理开销？

好问题！处理超大规模数据时，单个写入tf.train.Example确实会因为频繁的IO操作拖慢效率，其实TensorFlow有不少批量写入的方案，能帮你大幅降低开销，我给你整理几个实用的方法：

1. 攒批后一次性写入（最直接的方案） #

你可以先把多个tf.train.Example序列化后的字节串攒到一个列表里，等数量达标（比如1000个，可根据内存情况调整）再一次性写入文件，这样能大幅减少IO调用次数。示例代码如下：

writer = tf.python_io.TFRecordWriter(path)
batch_size = 1000
example_batch = []

for features in large_dataset:
    ex = generate_example(features)
    example_batch.append(ex.SerializeToString())
    
    # 攒够一批就写入
    if len(example_batch) >= batch_size:
        writer.write(b''.join(example_batch))
        example_batch = []

# 处理最后一批不足batch_size的数据
if example_batch:
    writer.write(b''.join(example_batch))
writer.close()

核心是用b''.join()把多个序列化字节串拼接成一个大的字节流，TFRecordWriter.write()本身支持写入包含多个Example的字节流，这样就把多次IO操作合并成了少数几次。

2. 借助tf.data.Dataset实现批量序列化与写入 #

如果你的数据已经是tf.data.Dataset格式，或者可以转成该格式，用Dataset的批量操作会更简洁，还能利用TensorFlow的并行处理能力加速序列化：

def serialize_example(features):
    # 复用你的generate_example逻辑，返回序列化后的字节串
    ex = generate_example(features)
    return ex.SerializeToString()

# 假设large_dataset是tf.data.Dataset对象
serialized_dataset = large_dataset.map(
    lambda x: tf.py_function(serialize_example, [x], tf.string),
    num_parallel_calls=tf.data.AUTOTUNE  # 并行序列化提升效率
)
batch_serialized_dataset = serialized_dataset.batch(batch_size=1000)

writer = tf.python_io.TFRecordWriter(path)
for batch in batch_serialized_dataset:
    # 拼接批量内的所有字节串并写入
    batch_bytes = tf.strings.join(batch).numpy()
    writer.write(batch_bytes)
writer.close()

搭配prefetch(tf.data.AUTOTUNE)还能让数据预处理和写入操作并行，进一步提升整体效率。

3. 多进程并行生成+批量写入（超大规模数据专属） #

如果数据量达到TB级别，单进程处理还是慢，可以用多进程解耦Example生成和写入操作：用多个进程并行生成序列化后的Example，放到队列里，再由一个专门的进程批量写入文件。示例代码如下：

import multiprocessing as mp

def writer_worker(features_queue, output_path, batch_size=1000):
    writer = tf.python_io.TFRecordWriter(output_path)
    batch = []
    while True:
        features = features_queue.get()
        if features is None:  # 收到结束信号
            if batch:
                writer.write(b''.join(batch))
            writer.close()
            break
        ex = generate_example(features)
        batch.append(ex.SerializeToString())
        if len(batch) >= batch_size:
            writer.write(b''.join(batch))
            batch = []

# 初始化队列和写入进程
queue = mp.Queue(maxsize=10)  # 队列大小根据内存调整
writer_process = mp.Process(target=writer_worker, args=(queue, path))
writer_process.start()

# 主进程负责给队列喂数据
for features in large_dataset:
    queue.put(features)

# 发送结束信号，等待写入进程完成
queue.put(None)
writer_process.join()

这种方式能最大化利用CPU资源，适合处理超大规模的数据集。

额外注意事项 #

• 批量大小要平衡内存占用和IO效率，一般建议在1000~10000之间调整；
• 无论用哪种方法，最后一定要关闭TFRecordWriter，避免数据丢失；
• 如果是分布式环境，还可以结合TensorFlow的分布式写入工具，但中小规模的超大规模数据，上面的方案足够应付。

内容的提问来源于stack exchange，提问作者bluesummers