数据管线：Dataset、DataLoader 与批处理

Dataset 与 DataLoader 各管什么

本节目标

本节围绕“Dataset 与 DataLoader 各管什么”展开。运行示例后记录样本索引、batch shape、padding mask 与数据划分，检查顺序和复现性。

清晰讲解

Dataset 定义“如何取得一条样本”，DataLoader 负责批处理、采样、并行加载和预取。拆开两者后，同一数据集可以替换不同采样策略而无需改读取逻辑。

核对“Dataset 与 DataLoader 各管什么”中的样本边界、batch 组成、mask、随机顺序和 train/validation 隔离。

代码示例

import torch
from torch.utils.data import Dataset, DataLoader

class PairDataset(Dataset):
    def __init__(self, n=100): self.x = torch.arange(n, dtype=torch.float32)
    def __len__(self): return len(self.x)
    def __getitem__(self, i): return self.x[i], 2*self.x[i] + 1

loader = DataLoader(PairDataset(), batch_size=8, shuffle=True)
print(next(iter(loader)))

运行结果与观察

一个 batch 应包含 8 个 x/y 对，且 y 始终满足 2*x+1；shuffle 只改变顺序不改变对应关系。

常见错误

• 在 __getitem__ 每次加载整个大文件；把随机划分逻辑散落在多个 worker 中。
• 排查“Dataset 与 DataLoader 各管什么”时只看 batch 能否产出，没有检查样本顺序、padding mask、重复与数据泄漏。

与大模型方向的连接

语料 Dataset 可以按文档、token 块或流式 shard 定义样本；DataLoader 决定每个 worker 与 rank 看到什么。

动手练习

为文本列表写一个 Dataset，返回文本长度和原始索引。

查看参考答案与验收点

固定 seed，打印首个 batch 的索引、shape 和 mask；再检查 train/validation 的稳定 hash 没有交集。

官方资料

• torch.utils.data API
• 数据加载教程

---

collate_fn 与变长序列 padding

本节目标

本节围绕“collate_fn 与变长序列 padding”展开。运行示例后记录样本索引、batch shape、padding mask 与数据划分，检查顺序和复现性。

清晰讲解

默认 collate 只能堆叠相同形状。变长文本需要在组 batch 时 padding，并同时返回有效位置 mask。按当前 batch 的最长序列补齐，比全局固定最大长度省计算。

核对“collate_fn 与变长序列 padding”中的样本边界、batch 组成、mask、随机顺序和 train/validation 隔离。

代码示例

import torch
from torch.utils.data import DataLoader

def collate(batch, pad_id=0):
    length = max(map(len, batch))
    ids = torch.full((len(batch), length), pad_id, dtype=torch.long)
    for i, seq in enumerate(batch): ids[i, :len(seq)] = torch.tensor(seq)
    return {'input_ids': ids, 'attention_mask': ids.ne(pad_id)}

print(collate([[2, 5, 8], [4, 9]]))

运行结果与观察

输出 input_ids 为 [2,3]，第二行末尾为 pad id，attention_mask 对应位置为 False。

常见错误

• pad id 同时是真实 token；只 padding 输入却不处理 labels；mask 的 dtype 或语义不符合模型接口。
• 排查“collate_fn 与变长序列 padding”时只看 batch 能否产出，没有检查样本顺序、padding mask、重复与数据泄漏。

与大模型方向的连接

LLM 的 input_ids、attention_mask 和 labels 必须在 padding、截断和 shift 之后保持对齐。

动手练习

扩展 collate：让 labels 的 padding 位置为 -100。

查看参考答案与验收点

固定 seed，打印首个 batch 的索引、shape 和 mask；再检查 train/validation 的稳定 hash 没有交集。

官方资料

• torch.utils.data API
• 数据加载教程

---

shuffle、Sampler 与数据顺序

本节目标

本节围绕“shuffle、Sampler 与数据顺序”展开。运行示例后记录样本索引、batch shape、padding mask 与数据划分，检查顺序和复现性。

清晰讲解

shuffle=True 适合普通训练；更复杂场景使用 Sampler 控制索引顺序。按长度分桶可减少 padding，但要保持足够随机性。验证集一般固定顺序，方便稳定对比。

核对“shuffle、Sampler 与数据顺序”中的样本边界、batch 组成、mask、随机顺序和 train/validation 隔离。

代码示例

import torch
from torch.utils.data import DataLoader, WeightedRandomSampler, TensorDataset
x = torch.randn(6, 3)
y = torch.tensor([0, 0, 0, 0, 1, 1])
weights = torch.where(y == 1, 0.5, 0.25)
sampler = WeightedRandomSampler(weights, num_samples=6, replacement=True)
loader = DataLoader(TensorDataset(x, y), batch_size=3, sampler=sampler)

运行结果与观察

采样 batch shape 保持正确，少数类出现频率在多次迭代后上升；单次随机结果不能证明分布。

常见错误

• 同时设置 sampler 和 shuffle；过采样后仍按原分布解释指标。
• 排查“shuffle、Sampler 与数据顺序”时只看 batch 能否产出，没有检查样本顺序、padding mask、重复与数据泄漏。

与大模型方向的连接

预训练数据常按 shard、语言和质量加权采样；采样策略本身就是训练配方的一部分。

动手练习

模拟长短文本并实现一个简单的按长度排序 batch sampler。

查看参考答案与验收点

固定 seed，打印首个 batch 的索引、shape 和 mask；再检查 train/validation 的稳定 hash 没有交集。

官方资料

• torch.utils.data API
• 数据加载教程

---

多进程加载、pin_memory 与 non_blocking

本节目标

本节围绕“多进程加载、pin_memory 与 non_blocking”展开。运行示例后记录样本索引、batch shape、padding mask 与数据划分，检查顺序和复现性。

清晰讲解

num_workers 让子进程并行准备数据，pin_memory 可帮助 CPU 到 CUDA 的异步拷贝，配合 .to(..., non_blocking=True) 使用。最佳 worker 数必须测量，不是越多越快。

核对“多进程加载、pin_memory 与 non_blocking”中的样本边界、batch 组成、mask、随机顺序和 train/validation 隔离。

代码示例

loader = DataLoader(
    dataset, batch_size=32, num_workers=4,
    pin_memory=torch.cuda.is_available(),
    persistent_workers=True,
)
for x, y in loader:
    x = x.to(device, non_blocking=True)
    y = y.to(device, non_blocking=True)

运行结果与观察

CPU 环境重点核对 batch 顺序与 shape；CUDA 环境再比较 pin_memory=True 配合 non_blocking=True 的拷贝时间，计时前后必须同步 GPU。

常见错误

• Windows 多进程脚本缺少 main guard；worker 太多造成内存和上下文切换压力。
• 排查“多进程加载、pin_memory 与 non_blocking”时只看 batch 能否产出，没有检查样本顺序、padding mask、重复与数据泄漏。

与大模型方向的连接

当 GPU 在等数据时，提高加载并行度或预取才有价值；若 tokenizer 是瓶颈，可先离线 tokenization。

动手练习

分别用 0、2、4 个 worker 测 100 个 batch 的耗时并记录环境。

查看参考答案与验收点

验收：首个 batch 的索引、shape 和 mask 可复现，并能说明“多进程加载、pin_memory 与 non_blocking”如何影响数据顺序与吞吐。

官方资料

• DataLoader 多进程与内存固定
• 数据加载教程

---

IterableDataset 与流式语料

本节目标

本节围绕“IterableDataset 与流式语料”展开。运行示例后记录样本索引、batch shape、padding mask 与数据划分，检查顺序和复现性。

清晰讲解

当数据巨大、来自网络流或不适合随机索引时，可用 IterableDataset 顺序产生样本。多 worker 时必须主动切分范围，否则每个 worker 会重复遍历全部数据。

核对“IterableDataset 与流式语料”中的样本边界、batch 组成、mask、随机顺序和 train/validation 隔离。

代码示例

from torch.utils.data import IterableDataset, get_worker_info

class RangeStream(IterableDataset):
    def __init__(self, end): self.end = end
    def __iter__(self):
        info = get_worker_info()
        start, step = (0, 1) if info is None else (info.id, info.num_workers)
        yield from range(start, self.end, step)

运行结果与观察

多 worker 汇总后应覆盖完整范围且没有重复；若每个 worker 都从 0 开始，则实现错误。

常见错误

• 多 worker 重复样本；流式数据没有可恢复游标；验证数据无限流导致评估不结束。
• 排查“IterableDataset 与流式语料”时只看 batch 能否产出，没有检查样本顺序、padding mask、重复与数据泄漏。

与大模型方向的连接

LLM 预训练常流式读取多个语料 shard；还要考虑 rank 与 worker 两层切分、断点位置和确定性洗牌。

动手练习

让示例同时按分布式 rank 和 worker id 切分且不重样。

查看参考答案与验收点

固定 seed，打印首个 batch 的索引、shape 和 mask；再检查 train/validation 的稳定 hash 没有交集。

官方资料

• torch.utils.data API
• 数据加载教程

---

数据划分、污染与可复现检查

本节目标

本节围绕“数据划分、污染与可复现检查”展开。运行示例后记录样本索引、batch shape、padding mask 与数据划分，检查顺序和复现性。

清晰讲解

训练、验证、测试的划分应在随机增强之前固定，并按任务实体选择合适粒度。文本任务只按行随机划分，可能让同一文档的近重复片段同时出现在训练和测试中。

核对“数据划分、污染与可复现检查”中的样本边界、batch 组成、mask、随机顺序和 train/validation 隔离。

代码示例

import torch
from torch.utils.data import random_split
g = torch.Generator().manual_seed(42)
train_set, val_set = random_split(dataset, [0.9, 0.1], generator=g)
print(len(train_set), len(val_set))

运行结果与观察

相同 generator 种子应产生完全相同的训练/验证索引；交集应为空、并集大小应等于原数据集，任何样本跨集合重复都视为污染。

常见错误

• 先切 chunk 再随机划分导致文档泄漏；只保存 seed 不保存实际数据版本和划分清单。
• 排查“数据划分、污染与可复现检查”时只看 batch 能否产出，没有检查样本顺序、padding mask、重复与数据泄漏。

与大模型方向的连接

大模型评测污染会制造虚高能力；去重和按来源隔离是数据工程的一部分，不是训练后的补救。

动手练习

给每条样本计算稳定 hash，检查 train 与 validation 是否存在交集。

查看参考答案与验收点

固定 seed，打印首个 batch 的索引、shape 和 mask；再检查 train/validation 的稳定 hash 没有交集。

官方资料

• random_split API
• 数据加载教程