模块 3 · Python、NumPy、PyTorch 与实验工程　/　第 7 课（共 10 课）

数据加载 Dataset 与 DataLoader

从零到前沿 ML 自学课程　·　阶段0：数学与工具基础　·　能力点：数据加载——Dataset/DataLoader/collate_fn 变长 padding，数据管线是隐形瓶颈

上一课你从零写出了完整训练循环：取一批数据 → 前向 → 算损失 → 反向 → 更新。但那时的"取一批数据"是手工切片的玩具数组。真实训练里，数据可能是几十万张图片、上百万条句子，长度还各不相同；它们要被高效、正确地源源不断喂进训练循环。本课就补上这块拼图：用 Dataset 描述"一条样本长什么样"，用 DataLoader 描述"怎么把样本攒成一个 batch"。数据管线是训练里最容易被忽视、却最常成为瓶颈的一环——GPU 再快，数据喂不上也是空转。

读完这一课，你将能够

  • 说清 Dataset 协议（__len__ 与 __getitem__）和 DataLoader 的四步流水线分别负责什么。
  • 从零手写一个可 shuffle、按 batch 切片的迭代器，逐 batch 打印 shape。
  • 写出处理变长序列的 collate_fn：右侧 pad 到同长，并生成 0/1 的注意力掩码（attention mask）。
  • 判断 transform / 数据增强该加在训练集还是验证集，并解释为什么。
  • 正确理解 num_workers 与 pin_memory，避免把它们当成"无脑调大就快"的旋钮。

Dataset：把"第 i 条样本"封装成一个对象

机制直觉。一个 Dataset 本质上只是一本"按编号取样本的字典"。它只需回答两个问题：一共有多少条（__len__）、第 i 条是什么（__getitem__(i)）。注意 __getitem__ 返回的是单条样本（比如一个 (特征, 标签) 元组），不是一批。攒成 batch 是 DataLoader 的活，分工要清楚。

这套"实现两个方法就能被框架使用"的约定叫协议（protocol，协议）。你不需要继承什么神奇基类，只要对象有 __len__ 和 __getitem__，DataLoader 就能用它——这正是 Python 鸭子类型的体现（更正式地说，这属于结构化类型 structural typing，本课不展开），和你上一课见到的 nn.Module 用 __call__ 的思路一脉相承。

DataLoader 四步流水线：从索引采样到 batch 张量 索引 → 取样本 → 拼接 → 张量　单向流动 ① 索引采样 索引池 [0,1,2,…,N-1] shuffle=True 每 epoch 打乱顺序 按 batch_size 取一批索引 2 0 这一批 = {2, 0} ② Dataset.__getitem__ dataset[2]　→ 样本 a 按索引返回一条样本 dataset[0]　→ 样本 b 逐条调用（单条） num_workers > 0 多个子进程并行执行这一步 proc 0 proc 1 proc 2 输出：样本列表 [样本 a, 样本 b]　(变长) ③ collate_fn 把变长样本列表 拼接 + 对齐 (padding) a 长 b 短 pad 补齐到统一长度 L ④ batch 张量 (B, L) 张量 + attention mask pin_memory 放入锁页内存 加速传往 GPU 送入模型 forward → 橙色虚线框 = 插在某一步上的"开关/配置"： shuffle 作用于 ①索引采样　·　num_workers 作用于 ②取样本（并行）　·　pin_memory 作用于 ④张量（搬运到 GPU 前）。它们不改变流向，只调节某一步的行为。
DataLoader 四步流水线：从索引采样到 batch 张量

最小手写。下面这个 ToyTextDataset 装的是已经被转成整数 id 的句子（真实 NLP 里这一步叫 tokenize，即分词，把文本切成 token id；本课先用现成 id）。注意每条句子长度不同——这正是后面 collate_fn 要解决的麻烦。

class ToyTextDataset:
    def __init__(self, sentences, labels):
        self.sentences, self.labels = sentences, labels

    def __len__(self):                 # 协议方法一：样本总数
        return len(self.sentences)

    def __getitem__(self, i):          # 协议方法二：返回第 i 条(单条)
        return self.sentences[i], self.labels[i]

sents  = [[5,7,9], [2], [4,4,4,8,1], [3,6]]   # 4 条变长句子(已转成 id)
labels = [1, 0, 1, 0]
ds = ToyTextDataset(sents, labels)

print("len(ds) =", len(ds))     # -> 4
print("ds[2]   =", ds[2])       # -> ([4, 4, 4, 8, 1], 1)  单条样本
print("ds[1]   =", ds[1])       # -> ([2], 0)

真实 API 当指针。PyTorch 里你继承 torch.utils.data.Dataset，实现同样的两个方法。基类几乎不做事，它的价值是给 DataLoader 一个统一的类型约定。

from torch.utils.data import Dataset

class TextDataset(Dataset):
    def __init__(self, sentences, labels):
        self.sentences, self.labels = sentences, labels

    def __len__(self):
        return len(self.sentences)

    def __getitem__(self, idx):
        return self.sentences[idx], self.labels[idx]

DataLoader：一条四步流水线

机制直觉。DataLoader 把零散样本变成成批张量，每一批走四步：

关键参数速记：batch_size（每批几条）、shuffle（是否每 epoch 打乱，训练开、验证关）、num_workers（几个子进程并行取数据，默认 0 即串行）、collate_fn（怎么拼 batch）、pin_memory（是否锁页加速传 GPU）。

最小手写。我们自己实现这条流水线的前三步。注意 shuffle 只打乱索引，不动原始数据；用一个带种子的随机源保证可复现（呼应本模块"可复现"主线）。collate_fn 先留个简单版本，下一节升级。

import random

class DataLoader:
    def __init__(self, dataset, batch_size, shuffle=False,
                 collate_fn=None, seed=0):
        self.ds, self.bs = dataset, batch_size
        self.shuffle = shuffle
        self.collate_fn = collate_fn
        self.rng = random.Random(seed)          # 可复现的打乱

    def __iter__(self):
        idx = list(range(len(self.ds)))         # 步骤1: 索引
        if self.shuffle:
            self.rng.shuffle(idx)               #         每 epoch 打乱索引
        for k in range(0, len(idx), self.bs):   # 步骤2: 按 batch_size 切片
            chunk = idx[k:k + self.bs]
            batch = [self.ds[i] for i in chunk]  #         逐个调 __getitem__
            yield self.collate_fn(batch) if self.collate_fn else batch

# 没有 collate_fn 时，直接吐出"样本列表"——还不是张量
loader = DataLoader(ds, batch_size=2, shuffle=False)
for b, batch in enumerate(loader):
    print(f"batch {b}: {batch}")
# batch 0: [([5, 7, 9], 1), ([2], 0)]      <- 两条长度不一，没法直接堆成矩阵!
# batch 1: [([4, 4, 4, 8, 1], 1), ([3, 6], 0)]

看最后两行注释：一批里两条样本长度不同，numpy 没法把 [5,7,9] 和 [2] 直接堆成一个规整矩阵。这正是默认 collate 遇到变长会直接报错的原因——下一节解决它。

collate_fn：把变长序列 pad 成等长 + 生成 mask【重点】

机制直觉。神经网络吃的是规整的张量，一个 batch 必须是形如 (B, L) 的矩阵，每行等长。但自然语言里句子天生不等长。解决办法是填充（padding，填充）：选一个特殊的 pad id（通常是 0），把这一批里较短的句子右侧补 0，补到和最长的那条一样长。

但这带来一个新问题：模型会把补出来的 0 当成真 token 去算注意力、算损失，污染结果。所以我们额外生成一个注意力掩码（attention mask，注意力掩码），形状和 batch 一样：真实 token 处标 1，padding 处标 0。模型靠它区分"哪些位置是真的、哪些是凑数的"。记住：pad 和 mask 是一对，永远成对出现。

变长序列右侧 padding 到同长，并生成 0/1 注意力掩码 原始：3 条不等长序列 本批最长 = 5 行1 5 7 9 行2 2 行3 4 4 4 8 1 右侧长度参差不齐 padding 后的 X　形状 (3, 5) 5 7 9 0 0 2 0 0 0 0 4 4 4 8 1 mask　形状 (3, 5) 1 1 1 0 0 1 0 0 0 0 1 1 1 1 1 逐格对齐 实色 token（真实） 浅灰斜纹 = padding (0) mask=1 mask=0 pad 与 mask 成对：mask=0 处即 padding，模型据此忽略凑数位置
变长序列右侧 padding 到同长，并生成 0/1 注意力掩码

一个常被忽视的细节：pad 的目标长度是这一批的最长句，不是全局最长句。所以不同 batch 的矩阵宽度可以不同（这叫 dynamic padding，比 pad 到全局最长省算力）。下面运行时你会看到不同 batch 的 X.shape 第二维确实不一样。

核心手写（完整可跑）。把 Dataset + DataLoader + collate_fn 三件套接起来，逐 batch 打印 shape 和 mask（其中 X[r, :len(s)] = s 用的就是模块 2 的切片赋值）：

import numpy as np, random

class ToyTextDataset:
    def __init__(self, sentences, labels):
        self.sentences, self.labels = sentences, labels
    def __len__(self):  return len(self.sentences)
    def __getitem__(self, i):  return self.sentences[i], self.labels[i]

def collate_fn(batch, pad_id=0):
    seqs   = [b[0] for b in batch]
    labels = [b[1] for b in batch]
    maxlen = max(len(s) for s in seqs)          # 只 pad 到"本批最长"
    X    = np.full((len(seqs), maxlen), pad_id, dtype=np.int64)
    mask = np.zeros((len(seqs), maxlen),        dtype=np.int64)
    for r, s in enumerate(seqs):
        X[r, :len(s)]    = s                    # 真实 token 写进去
        mask[r, :len(s)] = 1                    # 对应位置 mask 标 1
    return X, mask, np.array(labels, dtype=np.int64)

class DataLoader:
    def __init__(self, ds, batch_size, shuffle=False, collate_fn=None, seed=0):
        self.ds, self.bs, self.shuffle = ds, batch_size, shuffle
        self.collate_fn, self.rng = collate_fn, random.Random(seed)
    def __iter__(self):
        idx = list(range(len(self.ds)))
        if self.shuffle: self.rng.shuffle(idx)
        for k in range(0, len(idx), self.bs):
            batch = [self.ds[i] for i in idx[k:k+self.bs]]
            yield self.collate_fn(batch)

ds = ToyTextDataset([[5,7,9],[2],[4,4,4,8,1],[3,6]], [1,0,1,0])
loader = DataLoader(ds, batch_size=2, shuffle=True, collate_fn=collate_fn, seed=42)
for b, (X, mask, y) in enumerate(loader):
    print(f"batch {b}: X.shape={X.shape}, y={y.tolist()}")
    print(" X   =", X.tolist())
    print(" mask=", mask.tolist())

运行后你会看到两批的宽度不同（一批 pad 到 5，一批 pad 到 3），每条短句右侧是 0，对应 mask 也是 0——pad 与 mask 严丝合缝地对齐。

真实 API 当指针。PyTorch 的 DataLoader 接受一个 collate_fn 参数，签名就是"收一个样本列表，返回拼好的 batch"。变长序列常用 pad_sequence 来填充：

import torch
from torch.utils.data import DataLoader
from torch.nn.utils.rnn import pad_sequence

def collate_fn(batch):
    seqs   = [torch.tensor(s) for s, _ in batch]
    labels = torch.tensor([y for _, y in batch])
    X    = pad_sequence(seqs, batch_first=True, padding_value=0)  # (B, L)
    mask = (X != 0).long()                                        # 0 处为 padding
    return X, mask, labels

loader = DataLoader(ds, batch_size=2, shuffle=True,
                    num_workers=4, pin_memory=True,
                    collate_fn=collate_fn)
for X, mask, y in loader:
    ...  # X.shape == (B, L)，mask 与之同形

transform 与数据增强：只加在训练集

机制直觉。预处理可以分两类。一类是确定性的、保证输入格式正确的步骤（缩放到固定尺寸、按训练集均值方差归一化），训练集和验证集都要用，而且必须一致。另一类是随机数据增强（random augmentation，随机数据增强，如随机水平翻转、随机裁剪、加噪声），目的是人为制造多样性、防止模型死记硬背训练样本——它只加在训练集。

为什么验证集不能加随机增强？因为验证的使命是"在固定、干净的数据上量出模型真实水平"。如果每次验证时图片都随机翻转裁剪，你的验证指标会带上随机抖动，无法跨 epoch 比较，也就失去了判断"模型是否变好"的意义。验证集要的是可复现的、确定的预处理。

from torchvision import transforms

# 训练集：确定性预处理 + 随机增强
train_tf = transforms.Compose([
    transforms.RandomResizedCrop(224),     # 随机裁剪(增强)
    transforms.RandomHorizontalFlip(),     # 随机翻转(增强)
    transforms.ToTensor(),
    transforms.Normalize(MEAN, STD),       # 确定性
])
# 验证集：只有确定性预处理，没有任何随机
val_tf = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),            # 固定中心裁剪，非随机
    transforms.ToTensor(),
    transforms.Normalize(MEAN, STD),       # 与训练用同一套均值方差
])

数据管线是隐形瓶颈

机制直觉。训练每一步其实是两段接力：CPU 端准备下一批数据（读盘、解码、增强、collate），GPU 端计算当前批。如果数据准备比 GPU 计算还慢，GPU 就得干等——你买的算力在空转。这就是"隐形瓶颈"：loss 在降、代码没报错，但 GPU 利用率只有 30%。

两个对症旋钮，但都要正确理解，别迷信：

常见误解澄清：num_workers 管"并行读数据"，pin_memory 管"快速传 GPU"，二者解决的是不同环节，常配合使用但不可互相替代。下一课讲 GPU 时会把这条传输链路接完整。

调一调，观察现象

微任务 1：对比 shuffle=False 与同 seed 的 shuffle=True。预期现象：shuffle=False 时索引始终是 [0,1,2,3,4] 顺序切片，batch 0 永远是前两条、跨多次运行完全固定；shuffle=True 时顺序被打乱，但只要用同一个 seed，每次运行得到的打乱顺序也完全一样（这就是"可复现"）。下面四行打印，前两行（shuffle=False）相同、后两行（同 seed 的 shuffle=True）也相同——一个靠"不打乱"固定、一个靠"固定种子"固定。训练时若不 shuffle，同样顺序反复喂入，模型可能学到"数据排列顺序"这种虚假规律。

import random
class DS:
    def __init__(s, x): s.x = x
    def __len__(s): return len(s.x)
    def __getitem__(s, i): return s.x[i]
def loader(ds, bs, shuffle, seed=0):
    rng = random.Random(seed); idx = list(range(len(ds)))
    if shuffle: rng.shuffle(idx)
    for k in range(0, len(idx), bs):
        yield [ds[i] for i in idx[k:k+bs]]

ds = DS(['a','b','c','d','e'])
print("shuffle=False, 第1次:", [b for b in loader(ds, 2, False)])
print("shuffle=False, 第2次:", [b for b in loader(ds, 2, False)])  # 与上一行相同
print("shuffle=True , 第1次:", [b for b in loader(ds, 2, True, seed=0)])
print("shuffle=True , 第2次:", [b for b in loader(ds, 2, True, seed=0)])  # 同 seed，与上一行相同
# 思考：把第3、4行的 seed 改成不同值，shuffle=True 的两次还会一样吗？

微任务 2：把 collate_fn 里的 maxlen = max(...) 改成固定值（比如写死 maxlen = 4）。预期现象：长度为 5 的句子会被静默截断，而短句仍补到 4。为什么：pad 长度若小于本批最长句，就装不下真实 token；这说明"pad 到本批最长"不是随便选的——它是恰好容纳所有句子的最小宽度。注意下面这格主动用了 min() 防越界，所以只会丢弃超出部分、不会报错；如果你把 X[r, :n] = s[:n] 改成直接写 X[r, :len(s)] = s（不写 min），numpy 才会因为 shape 不匹配而报错——这一步留给你自己试。

import numpy as np
def collate(batch, pad_id=0, fixed=None):
    seqs = [b for b in batch]
    maxlen = fixed if fixed is not None else max(len(s) for s in seqs)
    X = np.full((len(seqs), maxlen), pad_id, dtype=np.int64)
    for r, s in enumerate(seqs):
        n = min(len(s), maxlen)        # 防越界：超出部分被丢弃
        X[r, :n] = s[:n]
    return X

batch = [[5,7,9], [2], [4,4,4,8,1]]
print("自动 maxlen:\n", collate(batch))            # 宽度=5，无丢失
print("写死 fixed=4:\n", collate(batch, fixed=4))   # 长度5的句子被静默截断!

微任务 3：把 batch_size 从 2 改成 3，观察最后一批。预期现象：bs=2 时 5 条样本切成 [2,2,1]、bs=3 时切成 [3,2]——两种 batch_size 演示的是同一件事：样本数不整除 batch_size 时，最后总会剩一个不满的"尾批/零头批"。为什么：样本数不一定整除 batch_size，默认保留这个零头批；真实 DataLoader 有 drop_last=True 可丢弃它（某些场景需要每批严格等大）。

def loader(n, bs, drop_last=False):
    idx = list(range(n)); out = []
    for k in range(0, n, bs):
        chunk = idx[k:k+bs]
        if drop_last and len(chunk) < bs: break
        out.append(chunk)
    return out

print("bs=2:", loader(5, 2))                  # [[0,1],[2,3],[4]]  尾批1条
print("bs=3:", loader(5, 3))                  # [[0,1,2],[3,4]]    尾批2条
print("bs=3, drop_last:", loader(5, 3, True)) # [[0,1,2]]          零头批被丢弃

动手练习

  1. 左 padding 版 collate。有些场景（如 decoder 生成）需要把 padding 补在左侧而非右侧。改写本课的 collate_fn，让短句左侧补 0，并相应生成 mask（真实 token 在右、mask 右侧为 1）。逐 batch 打印验证 mask 和 X 的 0 位置一致。
  2. 带 transform 的 Dataset。给 ToyTextDataset 加一个 transform 参数（一个函数），在 __getitem__ 里对样本调用它。再写一个"训练用"transform（给序列随机丢一个 token，模拟增强）和一个"验证用"transform（什么都不做）。验证：同一条样本，训练 transform 多次调用结果可能不同，验证 transform 永远相同。
  3. 测量 num_workers 的意义（用 sleep 模拟）。在纯 Python 里无法真用多进程展示，但你可以用 time.sleep(0.01) 模拟"取一条样本很慢"，然后串行遍历整个 dataset 计时。思考题（写注释回答）：如果有 4 个 worker 并行，理想情况下这段时间会变成原来的几分之几？为什么实际达不到这个理想值？

掌握自检