Dataset

Dataset 是一个 “样本索引 → 单条样本数据” 的映射
它只回答一个问题：给定第 i 个样本，应该返回什么？
Dataset 只负责三件事
- 定位样本
```
self.samples[i]
```
- 读取样本：从磁盘 / 内存 / 网络读取样本
- 返回模型可用的最小结构，如 Tensor / 数值 / 字典，返回的是单条样本，而不是 batch
```
return {
    "input": x,
    "target": y
}
```
Dataset 不应该做的事
- 不做 batch：batch 是 DataLoader 的职责
```
# 错误示例
def __getitem__(self, idx):
    return batch_x, batch_y
```
- 不做 shuffle：shuffle 是采样策略，应由 sampler 控制，多 worker 时不可控
```
random.shuffle(self.samples)
```
- 不放 GPU：Dataset 在 worker 进程中，而 GPU 只能在主进程 / 训练循环，会导致显存泄漏或 crash
```
x = x.to("cuda")   # 错误
```
- 不依赖全局可变状态：worker 是多进程，状态不同步，行为不可复现
```
global counter
counter += 1
```

Dataset 的组成结构

class MyDataset(torch.utils.data.Dataset):
    def __init__(self, data_list):
        self.data = data_list

    # Dataset 的逻辑长度，sampler / DataLoader 都会依赖它
    # 不等于 epoch 样本数，分布式会自动切分数据集
    def __len__(self):
        return len(self.data)

    # 给定 idx，返回一条确定的样本，Dataset 应该是幂等的
    def __getitem__(self, idx):
        item = self.data[idx]
        x = item["x"]
        y = item["y"]
        return x, y

Dataset 输出的常见风格
- Tuple：适合单任务
```
return x, y
```
- Dict：模型接口稳定，容易扩展
```
return {
    "input_ids": x,
    "labels": y,
    "attention_mask": mask
}
```
- NamedTuple /dataclass：更强的类型约束，适合大型工程

Map-style Dataset（标准 Dataset）

class MyDataset(torch.utils.data.Dataset):
    def __len__(self):
        return N

    def __getitem__(self, idx):
        return sample

对于 Map-style Dataset，因为能 index，PyTorch 可以：
- shuffle（通过 sampler）
- 切分数据（train /val）
- 分布式切 shard
- resume（按 index 恢复）
- random access augmentation
此时，Map-style Dataset 对应的 DataLoader 的行为模型如下
```
sampler → indices → __getitem__ → collate → batch
```

Iterable-style Dataset（IterableDataset）

class MyIterableDataset(torch.utils.data.IterableDataset):
    def __iter__(self):
        for sample in stream:
            yield sample

Iterable-style Dataset 只有 __iter__，没有 __len__，无法随机访问
- 这是由于有些数据没有第 i 条或者代价巨大
- 比如日志流、无限数据生成器、超大数据、Kafka / WebDataset
此时，Iterable-style Dataset 对应的 DataLoader 的行为模型如下
```
# 没有 sampler
__iter__ → sample → collate → batch
```

如果必须使用 Iterable-style Dataset，就必须自己承担责任

# 多 worker 下手动切分，否则每个 worker 读全量数据，会造成数据重复
def __iter__(self):
    worker_info = torch.utils.data.get_worker_info()

    if worker_info is None:
        start = 0
        step = 1
    else:
        start = worker_info.id
        step = worker_info.num_workers

    for i, sample in enumerate(self.stream()):
        if i % step == start:
            yield sample
# 分布式训练（数据并行）还需要再进行切分，需要额外考虑 rank、world_size

DataLoader

DataLoader 本质上是一个样本调度 + 批构造 + 并行加载的编排流程
DataLoader 不考虑数据语义，只做四件事：
- 决定 “按什么顺序” 取样本
- 决定 “一次取多少”
- 决定 “谁来取”（主进程 /worker）
- 决定 “怎么拼成 batch”

最常见的 Map-style Dataset 的迭代流程

Dataset
   ↓ (index)
Sampler / shuffle
   ↓ (indices)
BatchSampler
   ↓ (list of indices)
workers 调用 dataset[i]
   ↓ (list of samples)
collate_fn
   ↓
batch

DataLoader 示例

import torch
from torch.utils.data import Dataset

# 输入一个整数 x，标签是 y = 2x + 1
class LinearToyDataset(Dataset):
    def __init__(self, n):
        self.n = n

    def __len__(self):
        return self.n

    def __getitem__(self, idx):
        x = torch.tensor([idx], dtype=torch.float32)
        y = 2 * x + 1
        return {
            "x": x,
            "y": y
        }

from torch.utils.data import DataLoader

dataset = LinearToyDataset(n=10)

loader = DataLoader(
    dataset,
    batch_size=4, # 批次大小为 4
    shuffle=True, # RandomSampler
    num_workers=0 # 单进程
)

for batch in loader:
    print(batch)
    break
# {
#     "x": tensor([[3.],
#                  [7.],
#                  [1.],
#                  [9.]]),
#     "y": tensor([[ 7.],
#                  [15.],
#                  [ 3.],
#                  [19.]])
# }
#

如果指定 shuffle=True，DataLoader 内部将使用 RandomSampler，sampler 生成索引顺序，等价于
```
sampler = RandomSampler(dataset)
batch_sampler = BatchSampler(sampler, batch_size)
```
batch_size=N 指得是 sampler 生成 N 个 index，依次调用 N 次 dataset[i]，得到 N 个 sample，交给 collate_fn，然后 collate_fn 返回 batch
由此可见，sampler 决定顺序，batch_sampler 决定分组
常见 sampler
- SequentialSampler（验证 / 测试）
- RandomSampler（训练）
- DistributedSampler（分布式）
num_workers
- num_workers = 0：所有 dataset [i] 在主进程，容易调试，但读取数据集太慢
- num_workers > 0：fork /spawn 出 worker 进程，每个 worker 各自持有一份 Dataset，worker 并行调用 __getitem__
- 注意，Dataset 会被 pickle，因此不能依赖共享状态，不能在 Dataset 里存 socket /file handler（除非手动管理）
DataLoader 不该处理的事
- 不该理解样本结构
- 不该做 padding
- 不该做 tokenizer
- 不该做 augmentation
这些事情应该交由 Dataset 或 collate_fn 处理

collate_fn

在上面的例子中，DataLoader 取出数据之后，会得到 4 个 Dataset 返回的 dict

[
  {"x": tensor([3.]), "y": tensor([7.])},
  {"x": tensor([7.]), "y": tensor([15.])},
  ...
]

而默认的 collate_fn 会将其处理为

{
  "x": torch.stack([...], dim=0),
  "y": torch.stack([...], dim=0)
}

collate_fn 实际上是在保证输出结构是规整的，这也是为什么 Dataset 最好返回 dict
- tuple → zip + stack
- list → zip + stack
- dict → 按 key stack

在工程上可以用如下方式调用

for batch in loader:
    x = batch["x"]
    y = batch["y"]

    pred = model(x)
    loss = criterion(pred, y)

    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

那么如何自定义 collate_fn，实际上，collate_fn 接收的东西只有一个 batch 的 sample 列表
```
samples = [dataset[i] for i in indices]
batch = collate_fn(samples)
```
所以，Dataset → 定义单样本结构，而 collate_fn → 定义 batch 结构

最小自定义 collate_fn

def collate_fn(samples):
    return samples
# [
#   sample_0,
#   sample_1,
#   sample_2,
#   ...
# ]
#

示例：返回变长序列

class TextDataset(Dataset):
    def __getitem__(self, idx):
        length = torch.randint(3, 10, (1,)).item()
        tokens = torch.randint(1, 100, (length,))
        return {
            "input_ids": tokens,
            "length": length
        }

Padding 到 batch 内最大长度

def collate_fn(batch):
    input_ids = [item["input_ids"] for item in batch]
    lengths = torch.tensor([item["length"] for item in batch])

    max_len = max(len(x) for x in input_ids)

    padded = torch.zeros(len(input_ids), max_len, dtype=torch.long)

    for i, x in enumerate(input_ids):
        padded[i, :len(x)] = x

    attention_mask = (padded != 0).long()

    return {
        "input_ids": padded,
        "attention_mask": attention_mask,
        "lengths": lengths
    }

padding 一定在 collate_fn 中，而不是 Dataset，因为 Dataset 不知道 batch 的上下文，而 padding 是 batch 级别的操作
在 collate_fn 中不能使用 GPU 操作，因为 collate 是在 worker 进程中的，而 CUDA 不能跨进程安全使用，应该在训练循环中 .to(device)

sampler

Sampler 是一个 index 生成器：

for idx in sampler:
    yield idx

# DataLoader 内部
for indices in batch_sampler:
    samples = [dataset[i] for i in indices]
    batch = collate_fn(samples)
    yield batch

# shuffle=True
DataLoader(dataset, shuffle=True)
# 等价于
sampler = RandomSampler(dataset)

常见 sampler 对应关系
- shuffle=False → SequentialSampler
- shuffle=True → RandomSampler
- DDP → DistributedSampler

手动指定 sampler

sampler = torch.utils.data.RandomSampler(dataset)

loader = DataLoader(
    dataset,
    sampler=sampler,
    batch_size=32,
    shuffle=False  # 被忽略
)

DDP（Distributed Data Parallel）：在数据并行中，每张 GPU = 一个独立 Python 进程，每个进程都有一份 Dataset、一份 DataLoader、一份 Sampler
- 在 DDP 里，一个参与者 = 一个进程 = 一张 GPU
- world_size：参与分布式训练的进程总数
- rank 是这个进程在整个并行世界里的 ID

而 DistributedSampler 中，给定 dataset size = $N$ ，world_size = $W$ ，rank = $r$ ，它会把 index 划分为

rank 0: indices[0], indices[W], indices[2W], ...
rank 1: indices[1], indices[W+1], indices[2W+1], ...
...

每个样本只会被一个进程看到，所有进程加起来 = 完整 dataset

在 DDP 中，shuffle 必须由 DistributedSampler 自己做 shuffle

sampler = torch.utils.data.DistributedSampler(
    dataset,
    num_replicas=world_size,
    rank=rank,
    shuffle=True
)

loader = DataLoader(
    dataset,
    batch_size=batch_size,
    sampler=sampler,
    num_workers=4,
    pin_memory=True
)

sampler 在每个 epoch 中都要重新随机分配

sampler.set_epoch(epoch)
# 如果不设置 epoch，那么每个 epoch 的 shuffle 顺序是 完全一样的
# 这是因为 DistributedSampler 用 epoch 作为随机种子的一部分
# 如果不更新 epoch → 随机数固定
# 正确做法
for epoch in range(num_epochs):
    sampler.set_epoch(epoch)
    for batch in loader:
        ...

在 DDP 中，假设全局 batch 大小为 256，world_size 为 8，那么每个 DP 设备的 batch 应该是 256/8=32，每个进程的 sampler 只给 1/8 的样本，DataLoader 每次拿 32 个，梯度通过 AllReduce 汇总 —— DDP 不会自动帮你拆 batch

Dataset 与训练循环

一个标准的训练循环

# {
#     "inputs": {
#         "input_ids": ...,
#         "attention_mask": ...
#     },
#     "targets": {
#         "labels": ...
#     }
# }
#
def train_one_epoch(
    model,
    dataloader,
    optimizer,
    device
):
    model.train()

    for batch in dataloader:
        batch = move_to_device(batch, device)

        outputs = model(**batch["inputs"])
        loss = criterion(outputs,**batch["targets"])

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

性能问题

数据流程如下

Disk / Network
      ↓
Dataset.__getitem__()
      ↓
Worker Process (num_workers)
      ↓
[ Prefetch Queue ]
      ↓
collate_fn
      ↓
Pinned CPU Memory (pin_memory)
      ↓
GPU (async copy)

DataLoader 性能问题，本质只有一个：GPU 的计算速度比数据处理得快
num_workers
- num_workers = 同时跑 Dataset.__getitem__ 的子进程数量
- Dataset 里通常包含文件读取、解码、tokenize、augmentation，这些 CPU-bound / IO-bound 操作非常适合多进程
- 本地 SSD + 简单处理，num_workers 一般设置为 4-8
- 重预处理（tokenize /decode），num_workers 一般设置为 8-16
- 网络存储，num_workers 一般取决于带宽大小
- num_workers 不是越大越好，num_workers 越大，CPU 占用越高，导致 worker 队列阻塞，性能反而会下降
prefetch_factor
- 每个 worker 预先准备的 batch 数，默认 prefetch_factor = 2
- 实际总预取 batch 数：num_workers × prefetch_factor
- 如果 batch 构造慢，GPU 计算很快，那么 GPU 会等待 batch 的处理，而增加 prefetch 可以减少气泡，提高吞吐的稳定性
- 在 GPU utilization 抖动或训练时偶尔卡顿时，可以尝试改动 prefetch_factor
pin_memory
- 普通的 CPU 内存可以被操作系统 swap，而 pinned memory 物理页被锁定，不能被操作系统 swap
- 当指定 pin_memory=True 时，batch 会被拷贝到 pinned CPU memory，.to(device, non_blocking=True) 可以异步进行
- 因此一般要成对配合，否则 pin_memory 几乎没有收益
```
for batch in loader:
    batch = move_to_device(batch, device, non_blocking=True)
```
- 当 GPU 训练 /batch size 较大 / 数据复制的 IO 开销成为瓶颈时，可以考虑启用 pin_memory
- 当启用 pin_memory 时，一定要配合 non_blocking 使用
non_blocking=True
- 在 tensor.to(device, non_blocking=True) 中，当 tensor 在 pinned memory 且 CUDA steam 可用时，可以将复制与计算重叠，提高 GPU 的利用率
persistent_workers
- 默认情况下，每个 epoch 会先 fork workers，训练后再回收 workers
- 而 fork 的开销是比较昂贵的，通过保留 workers，可以使 workers 在 epoch 之间复用，大幅降低 epoch 切换开销
- 当 num_workers 大于 0，且 Dataset 的生命周期大于多个 epoch 时，可以启用该选项

性能友好的 DataLoader 配置

loader = DataLoader(
    dataset,
    batch_size=32,
    sampler=sampler,
    num_workers=8,
    pin_memory=True,
    prefetch_factor=4,
    persistent_workers=True
)

80% 的 GPU 空闲来自 IO、tokenize 等重预处理、Python 逻辑等，可以通过 GPU utilization、torch.profiler、Nsight Systems 等来判断训练缓慢的瓶颈

数据集