为了确保训练的模型能够复现,这里记录一下随机种子的设置,以便直接使用。
随机种子就像你玩游戏时的 "存档点",只要你加载同一个存档(设置同一个种子),后续的游戏进程(生成的随机数)就会完全一样;换个存档(换种子),进程就会不同。
与pytorch训练深度模型相关的随机种子主要在numpy,torch,random库上。下面直接上代码:
python
import random
import numpy as np
import os
import torch
def set_pytorch_seeds(seed=42, strict_reproducibility=False):
"""
设置PyTorch深度学习训练所需的所有随机种子,保证实验可复现
Args:
seed: 随机种子值,推荐使用固定整数(如42、123)
"""
# Python基础随机种子
random.seed(seed)
# NumPy随机种子
np.random.seed(seed)
# 禁用Python哈希随机化
os.environ['PYTHONHASHSEED'] = str(seed)
# PyTorch核心种子
torch.manual_seed(seed)
# GPU相关种子(单/多GPU)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
# 强制cudnn使用确定性算法
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# 仅在严格复现时启用以下配置
if strict_reproducibility:
# 禁用TF32(Ampere GPU生效)
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
# 配置CUBLAS + 确定性算法
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
torch.use_deterministic_algorithms(True)
# DataLoader种子配置(PyTorch多进程必备)
def get_pytorch_dataloader_config(seed: int = 42):
"""获取PyTorch DataLoader的种子配置,解决多进程随机性问题"""
def seed_worker(worker_id):
worker_seed = torch.initial_seed() % 2**32
np.random.seed(worker_seed)
random.seed(worker_seed)
generator = torch.Generator()
generator.manual_seed(seed)
return seed_worker, generator
- 核心函数set_pytorch_seeds仅保留 PyTorch 框架所需的所有关键种子配置,覆盖 CPU/GPU 场景,保证可复现性。
- 辅助函数get_pytorch_dataloader_config专门解决 PyTorch DataLoader 多进程的随机性问题,是完整复现的必要补充。