随机种子与 DataLoader：shuffle 和 sampler 的秘密

这一讲，我们聊聊一个让很多人头疼的问题——随机种子是怎么影响数据加载顺序的。

说实话，我之前带过好几个项目组，大家刚开始都会忽略这个细节。直到有一次，一个同学发现每次训练出来的模型指标波动特别大，排查了两天才发现——原来是 DataLoader 的 shuffle 没有固定种子。嗯，这里面的坑，我今天一次给你讲清楚。

DataLoader 的 shuffle 是怎么工作的？

先看最基础的情况。你用 PyTorch 的 DataLoader 加载数据，设置 shuffle=True。这时，框架会在每个 epoch 开始时，对数据集的索引进行一次随机排列。

这个随机排列的底层逻辑，靠的就是 Python 内置的随机数生成器，或者 NumPy 的随机状态。如果你不设置全局种子，那每次运行的结果都不一样。

看个最直接的例子：

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

# 创建一个简单数据集
data = TensorDataset(torch.arange(20))
loader = DataLoader(data, batch_size=5, shuffle=True)

# 第一次运行
print("第一次遍历结果：")
for batch in loader:
    print(batch[0].tolist())

print("\n第二次遍历结果（同一进程内）：")
for batch in loader:
    print(batch[0].tolist())

同一个进程内，每次遍历 DataLoader，shuffle 都会重新生成新的随机顺序。这很关键——同一个 epoch 内部，每个 batch 的顺序是确定的吗？ 不，恰恰相反，每次遍历都会重新洗牌。

Sampler：真正控制数据顺序的幕后角色

DataLoader 的 shuffle 参数，本质上是在帮你创建一个 RandomSampler。

我建议你记住这张对应表：

想自己控制种子？我给你两个办法：

方法一：设置全局种子

在脚本最开头调用：

torch.manual_seed(42)
# 如果用了 NumPy 操作
import numpy as np
np.random.seed(42)
# 如果用了 Python 随机
import random
random.seed(42)

方法二：显式传入 generator

这是更推荐的做法——精细控制，不污染全局状态：

g = torch.Generator()
g.manual_seed(42)

loader = DataLoader(dataset, batch_size=32, shuffle=True, generator=g)

for epoch in range(num_epochs):
    g = torch.Generator()
    g.manual_seed(42 + epoch)
    loader = DataLoader(dataset, batch_size=32, 
                        shuffle=True, generator=g)

多 Worker 情况下的种子传播

这里有个大坑——我当年在分布式训练时踩过。

当你设置 num_workers=4 时，DataLoader 会启动多个子进程来预取数据。每个子进程都会复制主进程的随机种子。如果你只设了主进程的种子，多个 worker 就可能生成相同的数据顺序。

PyTorch 的设计是：每个 worker 启动时，会调用 worker_init_fn 来重新设置种子，确保各 worker 的随机序列独立。默认实现是这样：

def _worker_init_fn(worker_id):
    torch.manual_seed(torch.initial_seed() + worker_id)
    np.random.seed(torch.initial_seed() + worker_id)
    random.seed(torch.initial_seed() + worker_id)

但是，如果你覆盖了 worker_init_fn，就一定要小心：

def my_worker_init_fn(worker_id):
    # 错误的做法：所有 worker 都用同一个种子
    # torch.manual_seed(42)  
    
    # 正确的做法：确保每个 worker 的种子不同
    torch.manual_seed(42 + worker_id)

loader = DataLoader(dataset, batch_size=32, 
                    num_workers=4,
                    worker_init_fn=my_worker_init_fn)

SubsetRandomSampler 的种子问题

我们在做数据集划分时，经常用 SubsetRandomSampler。它也需要关注种子：

from torch.utils.data import SubsetRandomSampler

# 划分训练集和验证集
indices = list(range(len(dataset)))
train_indices = indices[:800]
val_indices = indices[800:]

# 如果这里不设种子，每次划分结果不同
np.random.seed(42)
np.random.shuffle(indices)

train_sampler = SubsetRandomSampler(train_indices)
val_sampler = SubsetRandomSampler(val_indices)

train_loader = DataLoader(dataset, sampler=train_sampler, batch_size=32)
val_loader = DataLoader(dataset, sampler=val_sampler, batch_size=32)

看到了吗？SubsetRandomSampler 在初始化时，会进行一次内部洗牌。这也依赖全局随机状态。

你想保持严格的可复现性？我建议你这样做：

# 创建独立的 generator
g = torch.Generator()
g.manual_seed(42)

sampler = SubsetRandomSampler(indices, generator=g)
loader = DataLoader(dataset, sampler=sampler, batch_size=32)

知识结构一览

我把这一章的核心逻辑画了出来。你可以对照着看，理解 DataLoader 内部的数据流和种子影响点：

一个完整的实验对比

最后，我用一个实际例子帮你直观理解：同一个数据集，三种不同的种子设置，数据加载顺序到底差多少。

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

dataset = TensorDataset(torch.arange(10))

# 实验1：不设置种子
print("=== 实验1：无种子 ===")
loader1 = DataLoader(dataset, batch_size=3, shuffle=True)
for i, batch in enumerate(loader1):
    print(f"Batch {i}: {batch[0].tolist()}")

# 实验2：设置全局种子
print("\n=== 实验2：全局种子 42 ===")
torch.manual_seed(42)
loader2 = DataLoader(dataset, batch_size=3, shuffle=True)
for i, batch in enumerate(loader2):
    print(f"Batch {i}: {batch[0].tolist()}")

# 实验3：使用 generator 控制
print("\n=== 实验3：generator 种子 42 ===")
g = torch.Generator()
g.manual_seed(42)
loader3 = DataLoader(dataset, batch_size=3, shuffle=True, generator=g)
for i, batch in enumerate(loader3):
    print(f"Batch {i}: {batch[0].tolist()}")

好了，关于 DataLoader 的种子控制，核心要点就这些。记住 Sampler 是真正的主角，种子是控制它的钥匙。下回你的实验结果波动太大，先查查 DataLoader 的种子设了没。