Hugging Face Datasets streaming大数据集怎样流式读取而不一次载入内存怎样完成一个可复用的最小示例

我第一次处理远程大数据集时，真正卡住的不是模型，而是“还没开始训练，数据已经把下载目录和内存占满了”。这类任务可以先把目标改成流式读取：让数据在迭代时逐条到达，只处理当前样本或一个很小的窗口。

官方地址：https://huggingface.co/docs/datasets/stream

把 streaming=True 传给 load_dataset()，得到的是 IterableDataset。它适合顺序扫描、抽样、预处理和训练输入，不适合按下标随机取样或频繁回看最后一条数据。

先确认任务确实需要 streaming

如果数据集很大、只想先看少量样本，或者本地文件还不想转换成 Arrow，流式模式很合适。它的关键变化不是“把完整数据搬到更快的缓存”，而是把读取动作推迟到迭代发生的时刻。这样程序的内存主要由当前样本、处理缓冲和你主动保留的结果组成。

我会先问自己三个问题：是否只需要顺序遍历？是否能接受网络或文件读取的延迟？是否不需要通过索引反复访问同一行？三个答案大多为“是”，再进入下一步。

用 streaming=True 建立可复用的数据流

最小入口只需要指定数据集、分片和 streaming=True。下面的函数不把样本收集成一个大列表，而是边读边写 JSONL；take() 只用于演示如何限制本次读取的数量。

from pathlib import Path
import json

from datasets import load_dataset


def export_head(output_file: str, limit: int = 20) -> int:
    # streaming=True 返回 IterableDataset，读取发生在迭代时。
    stream = load_dataset(
        "HuggingFaceFW/fineweb",
        split="train",
        streaming=True,
    )

    written = 0
    # take() 限制示例数量，避免演示代码无意中扫完整个数据集。
    with Path(output_file).open("w", encoding="utf-8") as target:
        for row in stream.take(limit):
            # 只写需要的字段，避免把无关列继续留在下游文件中。
            record = {"text": row.get("text", ""), "token_count": row.get("token_count")}
            target.write(json.dumps(record, ensure_ascii=False) + "\n")
            written += 1
    return written


if __name__ == "__main__":
    # 真实任务可把 limit 改为批次上限，并按行交给后续处理器。
    print(export_head("fineweb-head.jsonl", limit=20))

这个例子只保留当前行和文件写入缓冲，结果文件会随着迭代增长，但不会因为数据集规模而在 Python 列表里积累全部样本。数据集标识、分片和字段名应根据实际数据卡调整。

把采样、跳过和分片放在正确的位置

如果需要近似随机抽样，可以在流上调用 shuffle(seed=42, buffer_size=10000)。这里的 buffer_size 是有限的随机缓冲，不等价于把整份数据打乱；缓冲越大，内存和读取开销也越高。

from datasets import load_dataset

# 先打乱，再切分；take/skip 会固定后续分片顺序。
stream = load_dataset("HuggingFaceFW/fineweb", split="train", streaming=True)
sampled = stream.shuffle(seed=42, buffer_size=10_000)
preview = sampled.take(100)
tail = sampled.skip(100)

# 逐条消费，避免 list(tail) 把剩余数据重新装入内存。
for row in preview:
    print(row.get("text", "")[:80])

数据由多个文件组成时，还可以使用 shard(num_shards, index) 切出一个分片；做并行训练时则要结合 worker 数量和分片数安排任务。需要反复遍历时，要明确每个 epoch 的读取代价，并按文档使用 set_epoch() 管理重新洗牌。

哪些任务不适合强行流式化

如果任务需要随机访问、按照全量数据排序、计算必须看到全部样本的精确统计，或者要在同一批数据上频繁回看，IterableDataset 会让代码变得别扭。官方说明也提醒，取最后一个样本可能需要先迭代前面的所有样本。

本地已有完整 Dataset 时，可以考虑 to_iterable_dataset()，它适合把本地数据转成可迭代形式，并可配置分片。反过来，如果只是想读取一个小数据集并多次索引，直接使用普通 Dataset 通常更直观。

结论与相关问题

这次方案的核心只有三点：用 streaming=True 延迟读取，用迭代器消费样本，用 take、shuffle、skip 或 shard 控制读取范围。它能降低一次性加载的压力，但不会消除网络延迟、单行过大或下游结果无限增长的问题；最终仍要控制自己的输出缓冲和写入策略。

常见追问：想只看前几条时用 take(n)；想抽样时先 shuffle 再 take；想多进程处理时先设计分片归属；想按索引随机访问时不要把 streaming 当成普通列表。