4、数据加载优化（下）：使用 TFRecord / WebDataset 加速 I/O，数据预读取与缓存策略

上一讲我们聊了数据加载的基本原理，说白了就是搞明白「数据到底卡在哪」。今天咱们来点硬核的——怎么用 TFRecord 和 WebDataset 这两个工具，把 I/O 瓶颈彻底打通。我个人习惯把数据加载优化分成两派：一派是「格式派」，另一派是「策略派」。今天两个都讲。

为什么小文件是性能杀手？

先问个问题：你训练集里有多少个小文件？几百个？几千个？还是几万个？

我在项目中遇到过最夸张的情况——一个同事的数据集有 50 万张图片，每张单独存成 JPEG。训练时 GPU 利用率只有 20%，剩下的时间全在等硬盘。为什么会这样？因为每次打开一个文件，操作系统都要做一次系统调用，寻道、读取、关闭。50 万次系统调用，你想想看，这得浪费多少时间。

说白了，小文件太多就是 I/O 的噩梦。解决方案也很直接：把小文件合并成大文件。

TFRecord：TensorFlow 的打包方案

TFRecord 是 TensorFlow 原生支持的二进制格式。它把多个样本打包成一个文件，顺序读取，效率极高。

来看个例子。假设你有一堆图片和标签，怎么转成 TFRecord？

import tensorflow as tf

def serialize_example(image_bytes, label):
    feature = {
        'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_bytes])),
        'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label]))
    }
    proto = tf.train.Example(features=tf.train.Features(feature=feature))
    return proto.SerializeToString()

# 写入 TFRecord
with tf.io.TFRecordWriter('train.tfrecord') as writer:
    for img_bytes, lbl in dataset:
        example = serialize_example(img_bytes, lbl)
        writer.write(example)

读取的时候，用 tf.data.TFRecordDataset 配合 map 解析就行：

def parse_fn(example_proto):
    feature_description = {
        'image': tf.io.FixedLenFeature([], tf.string),
        'label': tf.io.FixedLenFeature([], tf.int64)
    }
    return tf.io.parse_single_example(example_proto, feature_description)

dataset = tf.data.TFRecordDataset('train.tfrecord')
dataset = dataset.map(parse_fn, num_parallel_calls=tf.data.AUTOTUNE)

WebDataset：更灵活的替代方案

TFRecord 虽然好，但有个问题——它跟 TensorFlow 绑得太紧。如果你用 PyTorch 或者 JAX，用 TFRecord 就有点别扭了。这时候 WebDataset 就派上用场了。

WebDataset 的设计思路跟 TFRecord 类似，也是把多个样本打包成大文件。但它用的是更通用的 tar 格式。说白了，就是把一堆文件打成 tar 包，读取时按需解压。

# 安装：pip install webdataset
import webdataset as wds

# 创建数据集
dataset = wds.WebDataset("shards/train-{000000..000009}.tar")

# 解码图片
dataset = dataset.decode("rgb").to_tuple("jpg", "cls")

# 构建 DataLoader
dataloader = torch.utils.data.DataLoader(
    dataset, batch_size=64, num_workers=4
)

WebDataset 的好处是：

• 框架无关：PyTorch、TensorFlow、JAX 都能用
• 支持流式读取：可以从网络、S3、本地文件系统读取
• 天然支持 sharding：多机多卡训练时自动分配 shard

数据预读取：让 GPU 永远不饿

格式优化完了，接下来是策略优化。数据预读取，说白了就是「提前把下一批数据准备好」。GPU 在算当前 batch 的时候，CPU 已经在加载下一个 batch 了。

在 tf.data 里，用 prefetch 就能实现：

dataset = dataset.batch(64).prefetch(tf.data.AUTOTUNE)

就这么一行代码，训练吞吐能提升 20%-50%。为什么？因为 GPU 和 CPU 可以并行工作，不再互相等待。

在 PyTorch 里，DataLoader 的 num_workers 参数就是干这个的：

dataloader = DataLoader(dataset, batch_size=64, num_workers=4, prefetch_factor=2)

prefetch_factor=2 表示每个 worker 预加载 2 个 batch。我一般设成 2 或 3，太大反而会占用过多内存。

缓存策略：内存换时间

如果数据集不大（比如几万张图片），完全可以全部缓存到内存里。这样训练时 I/O 延迟几乎为零。

# TensorFlow 缓存
dataset = dataset.cache(filename='cache_dir')  # 也可以不指定文件名，直接缓存到内存

# PyTorch 缓存
class CachedDataset(Dataset):
    def __init__(self, dataset):
        self.data = [dataset[i] for i in range(len(dataset))]
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]

但要注意，缓存不是万能的。如果数据集太大，缓存反而会触发 swap，拖慢整个系统。我建议先看看内存占用，留出 20% 的余量给模型参数和中间变量。

一张图看懂数据加载优化

下面这张图总结了今天讲的核心内容：

实战建议：先诊断，再优化

别一上来就改格式、加缓存。先看看你的 I/O 到底卡在哪。我一般用 nvidia-smi 看 GPU 利用率，用 iostat 看磁盘 I/O。如果 GPU 利用率低于 50%，磁盘 I/O 接近 100%，那基本就是数据加载的问题。

优化顺序我建议这样：

1. 先加 prefetch：一行代码，立竿见影
2. 再调 num_workers：从 2 开始，逐步增加到 8
3. 如果还不够，考虑格式优化：转成 TFRecord 或 WebDataset
4. 最后上缓存：数据集小的话直接全量缓存

嗯，数据加载优化这块基本就这些了。记住一个原则：别让 GPU 闲着。只要 GPU 利用率上去了，训练吞吐自然就快了。

想要更多资料，联系可可