为什么学这个
最近在做图像超分辨率算法的加速优化研究,需要进行大量的消融实验来验证不同模块(比如残差块层数、注意力机制等)的有效性。
在这个过程中,我遇到了一个极其让人崩溃的问题:明明在代码里固定了所有的随机数种子(设为123),但每次跑出来的首个 Epoch 的验证集 PSNR 依然有细微差异。 对于严谨的算法工程来说,如果基线(Baseline)无法完美复现,所谓的“指标提升”就成了无源之水。为了彻底弄清底层的随机性机制,我把 PyTorch 的训练流底朝天翻了一遍,终于理清了实现“绝对确定性”的标准打法,在此做个总结记录。
核心内容/步骤
要做到真正的“完美复现”,仅仅写一句 torch.manual_seed() 是远远不够的。必须从四个维度彻底封死随机性:
-
系统与基础库环境:消除 Python 字典哈希的随机性,并固定
random和numpy的种子。- Python Hash 随机化:设置环境变量
PYTHONHASHSEED = str(seed)。 - Python 原生 random 库:
random.seed(seed)。 - NumPy 库:
np.random.seed(seed)。
- Python Hash 随机化:设置环境变量
-
PyTorch 与 GPU 环境:固定 CPU 和所有 GPU 卡的计算种子。
- CPU 种子:
torch.manual_seed(seed)。 - GPU 种子:
torch.cuda.manual_seed(seed)以及torch.cuda.manual_seed_all(seed)(针对多卡)。
- CPU 种子:
-
底层计算加速库(cuBLAS/cuDNN) :现在的显卡为了追求极致算力,默认会使用带有随机性的原子操作或动态选择卷积算法。必须强制关闭 cuDNN 的自动寻优,并开启确定性算法。
- 关闭自动寻优:
torch.backends.cudnn.benchmark = False。 - 开启确定性卷积:
torch.backends.cudnn.deterministic = True。 - 配置 cuBLAS 工作空间:设置环境变量
CUBLAS_WORKSPACE_CONFIG = ':4096:8'或':16:8'。 - 终极杀器:调用
torch.use_deterministic_algorithms(True)。开启后,如果代码中用到了无法保证确定性的算子,PyTorch 会直接报错,而不是默默产生随机结果。
- 关闭自动寻优:
-
数据加载(DataLoader) :多进程数据加载是破坏复现的重灾区,必须为每个 worker 分配独立的派生种子,并绑定固定的 Generator 控制 shuffle 行为。
- 生成独立的 Generator:
g = torch.Generator(); g.manual_seed(seed),并传给 DataLoader。 - 初始化子进程种子:编写
worker_init_fn,利用torch.initial_seed()为每个 worker 分配固定的派生种子。
- 生成独立的 Generator:
注意事项:PyTorch 复现的高频“踩坑”点
即使你对齐了训练策略并锁死了随机种子,依然有可能在以下几个隐蔽的角落翻车:
-
坑位 1:环境变量设置时机不对 如果你使用
CUBLAS_WORKSPACE_CONFIG来约束 CUDA 的底层算子,该环境变量必须在任何 CUDA 操作被调用之前(通常在import torch之后的第一行)设置。如果在张量初始化后设置,将彻底失效。 -
坑位 2:上采样/插值算子的非确定性 在超分辨率或分割任务中,部分插值算法(如
F.interpolate)在 GPU 上的反向传播目前没有确定性实现。如果开启了强制确定性算法,程序会直接崩溃。对策是寻找替代算子或将相关操作转移到 CPU。 -
坑位 3:第三方图像增强库的“随机黑洞” 除了
torchvision,如果你使用了Albumentations、imgaug或OpenCV,它们内部有独立的随机状态。必须在 DataLoader 的worker_init_fn中为它们单独设置种子。 -
坑位 4:实验顺序污染 在同一个 Python 脚本中用
for循环连续跑多个消融实验时,前一个实验会消耗全局随机数生成器的状态。必须在每次独立实验开始前,重新调用一次全局set_seed(),确保每个实验起点一致。 -
坑位 5:忘记切换 Eval 模式 测试时如果忘记调用
model.eval(),Dropout 或 BatchNorm 层会继续按照训练模式引入随机性,导致每次验证集输出都不一样。
收获与总结
深度学习里的“随机性”就像是调皮的幽灵。完美复现不仅仅是一个代码规范问题,它是算法工程师排查 Bug、验证猜想的基石。
这次踩坑让我深刻认识到:控制变量法是做科研和工程的核心底线。无论是数据增强的第三方库、模型的 Dropout 层,还是哪怕少写了一行梯度裁剪,都会在深度神经网络的非线性放大下,导致最终结果的蝴蝶效应。严谨,才是算法落地的第一生产力。
附录:实现完美复现的标准步骤和注意事项代码
在你的项目根目录下,建议封装以下两个核心函数,并在运行任何模型代码前第一时间调用。
Python
import os
import random
import numpy as np
import torch
import torch.backends.cudnn as cudnn
def set_seed(seed=123):
"""
设置所有随机源种子,开启 PyTorch 终极绝对确定性模式
注意:此函数需要在导入其他第三方运算库之前尽早调用
"""
# 1. 消除 Python 内置数据结构的哈希随机性
os.environ['PYTHONHASHSEED'] = str(seed)
# 2. 强制 cuBLAS 使用确定性工作空间 (必须在 CUDA 初始化前设置)
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
# 3. 固定常规随机数生成器
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed) # 针对多卡
# 4. 约束 cuDNN 行为
cudnn.deterministic = True # 强制使用确定性卷积算法
cudnn.benchmark = False # 关闭自动寻优
# 5. 【终极杀器】强制全局算子必须使用确定性算法!
# 开启后,如果遇到无法确定性的算子(如某些插值操作的反向传播),程序会直接报错而不是默默产生随机结果
torch.use_deterministic_algorithms(True)
def worker_init_fn(worker_id):
"""
DataLoader worker 初始化函数,防止多进程数据加载产生相同的随机序列
"""
worker_seed = torch.initial_seed() % 2**32
np.random.seed(worker_seed)
random.seed(worker_seed)
# 如果使用了 imgaug 或 albumentations,这里也需要对它们进行单独的 seed 设置
# import imgaug; imgaug.random.seed(worker_seed)
# =========================================
# 使用示例:DataLoader 的正确挂载方式
# =========================================
'''
g = torch.Generator()
g.manual_seed(123)
train_loader = DataLoader(
dataset,
batch_size=16,
shuffle=True, # 开启 shuffle
num_workers=4,
worker_init_fn=worker_init_fn, # 必须传入初始化函数
generator=g # 绑定生成器控制 shuffle 的确定性
)
'''