Transformers datasets流式读取训练数据

训练集大到无法完整放进本地磁盘时，Hugging Face Datasets 可以用 streaming=True 按迭代器逐步读取。它返回的不是支持随机下标和完整长度的普通 Dataset，而是 IterableDataset：数据在遍历时才被取出，适合直接做预览、预处理和按训练步数消费。

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

最小方案是：load_dataset(..., streaming=True) 创建流式数据集，需要随机化时先调用 shuffle(buffer_size=...)，再用 map 做处理；交给 Transformers 训练时不要依赖 len(dataset)，用 max_steps 明确训练边界。

streaming=True 为什么适合大训练集

普通加载方式通常会准备一个可索引的数据集对象；流式加载则把重点放在“边迭代边读取”。当数据集体积超过磁盘容量、只想先取少量样本，或者训练任务本来就按批次持续消费时，这种方式可以减少启动前的完整下载等待。

调用 load_dataset 时只需要把 streaming 设为 True。返回对象的类型变化很重要：它表现为可迭代的数据流，而不是可以随时跳到第 n 条记录的数组。

最小写法：先流式取几条样本

先不要把整条训练流水线搬进来，可以用 take 取一个小窗口，确认字段名称和数据形状。这里的 take(3) 表示最多消费前三条，不会为了计算总数而扫描完数据集。

from datasets import load_dataset

# streaming=True 让样本在迭代时读取，不先完整下载训练集。
dataset = load_dataset(
    "HuggingFaceFW/fineweb",
    split="train",
    streaming=True,
)

# take 只取一个小窗口，适合先检查字段和内容形状。
for index, row in enumerate(dataset.take(3)):
    # 只打印少量字段，避免把整条样本输出到日志。
    print(index, row.get("text", "")[:120])

如果数据来自本地文件，也可以把 data_files 指向文件或通配路径，再保持 streaming=True。流式并不限定数据必须来自 Hub，关键在于读取方式和返回对象的迭代语义。

IterableDataset 和普通 Dataset 的差别

因此，下面这类写法在流式模式下不应作为主流程：

# 不要把 IterableDataset 当成可以随机下标访问的列表。
first_row = dataset[0]

# 不要依赖总长度来决定一轮训练的步数。
total_rows = len(dataset)

如果业务必须随机访问、精确计算总条数或反复从任意位置读取，先使用普通 Dataset 或准备合适的本地分片会更直接。流式模式解决的是“边读边用”，不是把所有数据集能力都换成惰性版本。

shuffle、map、take 的顺序要固定

流式数据集可以用 shuffle(buffer_size=...) 做有限缓冲区随机化。缓冲区只保存指定数量的候选样本，随机化质量和内存占用都会受到 buffer_size 影响；它不是完整数据集级别的全量洗牌。

from datasets import load_dataset

# 创建流式数据源，长度可以未知。
stream = load_dataset(
    "HuggingFaceFW/fineweb",
    split="train",
    streaming=True,
)

# 先打乱分片和有限缓冲区，再截取本轮要消费的窗口。
shuffled = stream.shuffle(seed=42, buffer_size=10_000)
window = shuffled.take(1_000)

for row in window:
    # 这里只演示消费边界；真实训练中可在 map 中完成 tokenize。
    text = row["text"]
    if text:
        print(text[:80])

take 和 skip 会锁定后续分片顺序，所以需要随机化时应先 shuffle 再切分。反过来先 take 再 shuffle，只会在已经截取的小窗口内进行随机化，不能达到“从全体流式数据中抽取”的效果。

预处理可以用 map 串在流上。每条样本在被消费时处理，适合轻量、确定性的字段转换；如果转换本身很重，要额外评估 CPU、网络读取和训练 GPU 的节奏是否匹配。

接入 Transformers 训练时用 max_steps 设边界

流式数据通常没有可直接依赖的总长度，因此训练循环不应以“完整遍历一遍数据集”作为唯一停止条件。使用 Transformers 的 Trainer 或类似训练器时，常见做法是指定 max_steps，让训练以优化步数结束。

from datasets import load_dataset
from transformers import AutoTokenizer, DataCollatorWithPadding, TrainingArguments

model_name = "distilbert/distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 训练数据按需读取，避免先把全部文本落到本地。
stream = load_dataset(
    "imdb",
    split="train",
    streaming=True,
)

def tokenize_batch(batch):
    # truncation 给单条文本设置上限，避免异常长样本拖垮批次。
    return tokenizer(batch["text"], truncation=True, max_length=256)

# batched 处理要与下游训练器的批次约定一致。
tokenized_stream = stream.map(tokenize_batch, batched=True)
collator = DataCollatorWithPadding(tokenizer=tokenizer)

args = TrainingArguments(
    output_dir="./outputs",
    per_device_train_batch_size=8,
    # 数据长度未知时，用明确步数代替按 epoch 推断总样本数。
    max_steps=1_000,
    logging_steps=50,
)

# 这里将 tokenized_stream 交给 Trainer 的 train_dataset 参数。
# 具体 Trainer、模型和标签列需按任务补齐，本文只展示流式边界。
print(args.max_steps, collator)
print(next(iter(tokenized_stream)))

上例把“数据能持续产生”与“训练何时结束”拆成两个问题：前者由 IterableDataset 和迭代器负责，后者由 max_steps 负责。若仍希望按 epoch 表达，需要先明确可用的样本数或分片计划，不能假设流式对象会自动提供精确长度。

并行、分片与随机访问的边界

如果从一个已经存在的普通数据集转换为 IterableDataset，可以通过分片让多个 DataLoader worker 分担读取。官方文档给出的思路是准备多个 shard，再使用 to_iterable_dataset(num_shards=...)，这样每个 worker 负责其中一部分，而不是让所有 worker 反复读同一份数据。

from datasets import load_dataset
from torch.utils.data import DataLoader

# 先得到可索引数据集，再转换成带分片能力的可迭代数据集。
materialized = load_dataset("ethz/food101", split="train")
iterable = materialized.to_iterable_dataset(num_shards=64)

# 缓冲区随机化只在迭代过程中生效，worker 数量应与分片数量匹配。
iterable = iterable.shuffle(seed=7, buffer_size=10_000)
dataloader = DataLoader(iterable, num_workers=4)

for batch in dataloader:
    # 真实项目里这里接模型训练；示例只消费一个批次。
    print(batch.keys())
    break

这条路线的代价是：你先准备了普通数据集，才获得了更方便的分片转换；它与“完全不下载远程大数据”的 load_dataset(..., streaming=True) 不是同一条路径。选择哪一种，取决于数据是否已经在本地以及是否需要 worker 级分片。

性能和安全注意

• 缓冲区不是全量洗牌。 增大 buffer_size 通常会增加内存和预取压力，先根据样本大小和 worker 数量估算。
• 不要把未知长度写成伪长度。 训练步数、梯度累积和 batch size 要一起换算，避免最后的训练量与预期不一致。
• 谨慎打印原始样本。 训练数据可能含个人信息、访问令牌或业务文本，预览时只打印必要字段和截断片段。
• 先 shuffle 再 take/skip。 需要随机采样时，这个顺序比把一个已经固定的小窗口再次随机化更符合预期。
• 区分网络瓶颈和模型瓶颈。 流式读取让下载与训练交错进行，网络、解压、tokenize 或 worker 任一环节变慢，都可能让 GPU 等待。

速查表：什么时候选流式读取

相关问题

streaming=True 会不会把数据集缓存成完整文件？

它的目标是迭代时读取，不要求先把整个数据集准备成完整本地数据集；具体缓存和底层文件行为仍取决于数据格式与加载器。

IterableDataset 能不能直接调用 len？

不要把它当成一定有长度的对象。流式数据的总样本数可能未知，训练时使用明确的 max_steps 更稳妥。

为什么 take 要放在 shuffle 后面？

take 会固定后续分片顺序。若先 take，后续 shuffle 只作用于已经截取的小窗口；先 shuffle 再 take 才能让窗口从流式随机化结果中产生。

流式数据能不能做 tokenize？

可以使用 map 在迭代过程中处理，常见做法是配合 batched=True 交给 tokenizer；需要关注 tokenizer 速度是否跟得上读取和训练。