广告位 · 页头(后台-模块管理-新建模块后替换此区域)

PyTorch重现性为什么总不稳定?随机种子与DataLoader的正确操作有哪些?

在 PyTorch 中追求可复现的实验结果,很多开发者会下意识地只调用一次 torch.manual_seed,但随后发现同样的代码在不同运行中仍然会产生不同的损失曲线或指标。这并非框架缺陷,而是深度学习训练链路中同时存在多个随机源。本文从随机数生成机制入手,梳理固定种子、数据加载、卷积计算和评估模式的完整操作要点,帮助你快速建立一套可落地的复现流程。

PyTorch重现性为什么总不稳定?随机种子与DataLoader的正确操作有哪些?

一、结果不稳定的随机性来源有哪些?

PyTorch 训练涉及 Python 内置 random、NumPy 随机数生成器、PyTorch CPU 与 GPU 随机数生成器、CUDA 卷积库 cuDNN 的算法选择,以及 DataLoader 多进程采样。这些随机源彼此独立,只要有一个没有被固定,训练结果就可能发生变化。Python 的 random 模块常用于数据划分与增强,NumPy 常用于数据预处理,PyTorch 张量初始化、Dropout、参数初始化等依赖自身 RNG。即使固定了 PyTorch,未固定 NumPy 时,数据增强结果仍会漂移,最终影响梯度更新。

cuDNN 是更隐蔽的因素。GPU 上的卷积、池化等操作默认会从多种算法中自动选择性能最优的实现,这些算法可能在非确定性计算顺序下产生极小浮点差异。训练初期差异很小,但随着迭代累积会放大为完全不同的权重轨迹。DataLoader 多进程读取数据时,每个 worker 会继承主进程的随机状态,但不同 worker 的调用顺序可能不同,shuffle 结果也会不同。

因此,想要实现 PyTorch 重现性,不能只关心 torch.manual_seed,需要建立一份涵盖 CPU、GPU、数据加载和算法选择的全量种子设置。下面给出可直接复用的操作代码。

二、固定随机种子的正确操作与代码模板

第一份种子设置模板应当同时覆盖 Python、NumPy、PyTorch CPU 和所有 GPU。下面这段代码可以放在训练脚本的入口处。

import random
import numpy as np
import torch
def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.use_deterministic_algorithms(True)

需要说明每一行的作用。random.seed 保证 Python 层面的数据划分、文件顺序等稳定;np.random.seed 固定 NumPy 随机数;torch.manual_seed 固定 PyTorch CPU 张量初始化和 Dropout 等随机操作;torch.cuda.manual_seed 和 manual_seed_all 用于当前 GPU 和多 GPU;cudnn.deterministic 为 True 后,cuDNN 只使用确定性算法,可能略微降低速度;cudnn.benchmark 必须为 False,否则自动算法搜索会引入随机性。torch.use_deterministic_algorithms(True) 是更严格的全局限定,但某些模型可能因此报错,需要根据场景开启。

要注意种子设置的调用时机。固定种子必须在任何随机数被消耗之前执行,包括模型初始化、数据集划分、数据增强等。如果某些第三方库内部使用了自己的随机源,还需要单独设置。例如 torchvision 的图像变换通常依赖 NumPy 或 PyTorch 随机,但自定义数据集如果使用 Python random,前面的 random.seed 就足以覆盖。

DataLoader 的 shuffle 在多 worker 下需要额外处理。即使主进程设置了种子,每个 worker 启动时会从主进程继承状态,但为了结果稳定,建议显式指定 generator 和 worker_init_fn。

from torch.utils.data import DataLoader
def worker_init_fn(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)
loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    worker_init_fn=worker_init_fn,
    generator=torch.Generator().manual_seed(42)
)

这里的 generator 参数会让 DataLoader 使用独立的随机生成器,不受全局种子改变影响,便于在多次实验中保持数据采样一致。worker_init_fn 则保证每个 worker 的 Python 与 NumPy 随机源都从确定值开始。

三、常见疑问解答:为什么固定种子后仍不稳定?

常见疑问一:Dropout 与 BatchNorm 是否会影响复现?Dropout 使用 PyTorch 随机数,固定种子后在相同网络结构下是稳定的。BatchNorm 在训练时使用当前 batch 的均值与方差,因此 batch 顺序和 batch 大小都会影响统计量;复现时需要保证数据顺序一致。测试时 BatchNorm 使用全局统计量,必须调用 model.eval(),否则继续更新统计量会导致明显的评估差异。

model.eval()
with torch.no_grad():
    predictions = model(inputs)

常见疑问二:单 GPU 能复现,多 GPU 为什么不行?多 GPU 训练中,数据并行会将一个 batch 拆分成多个子 batch,不同 GPU 上的 BatchNorm 会分别计算局部统计量,之后同步全局均值与方差。如果不同卡的计算顺序或通信顺序不同,浮点累加顺序也会变化,导致微小差异。可以采用单 GPU 复现验证、使用同步 BatchNorm 替代普通 BatchNorm,或者固定设备顺序与数据划分来缓解。

常见疑问三:torch.use_deterministic_algorithms(True) 和 cudnn.deterministic 有什么区别?前者是 PyTorch 层面的强制确定性,能捕获更多非确定性算子,但兼容性更严格;后者只作用于 cuDNN 后端。两者可以同时开启,但若模型包含不支持确定性的操作,训练会直接报错,需要针对性关闭。

四、少走弯路的检查清单

排查复现问题时,可以按照从数据到模型再到环境的顺序检查。首先是数据:原始数据集是否一致、划分方式是否随机且已固定、DataLoader 是否设置了 generator 与 worker_init_fn。其次是模型:权重初始化是否在种子设置之后、Dropout 是否在同一位置、BatchNorm 统计量是否受 batch 顺序影响。最后是环境:PyTorch 版本与 CUDA 版本是否有变化、是否开启了 cudnn.benchmark、是否在 CPU 与 GPU 之间切换。

另外,保存与加载模型权重时,应使用 state_dict 而不是整个模型对象,可以避免 Python 版本或类定义导致的额外不确定因素。评估阶段使用 torch.no_grad() 关闭梯度计算,同时调用 model.eval() 固定 BatchNorm 与 Dropout 行为,这也是获得稳定指标的前提。

最后提醒,完全复现不应当以牺牲过多训练速度为代价。对于要求严格的论文实验,可以开启全部确定性选项;对于一般调参实验,固定预设种子并保持数据顺序一致通常足够。关键是根据自身需求平衡确定性与运行效率。

广告位 · 文章底部(后台-模块管理替换此区域)