pytorch深度学习训练随机种子设置

为了确保训练的模型能够复现,这里记录一下随机种子的设置,以便直接使用。

随机种子就像你玩游戏时的 "存档点",只要你加载同一个存档(设置同一个种子),后续的游戏进程(生成的随机数)就会完全一样;换个存档(换种子),进程就会不同。

与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 多进程的随机性问题,是完整复现的必要补充。
相关推荐
工程师老罗10 小时前
举例说明YOLOv1 输出坐标到原图像素的映射关系
人工智能·yolo·计算机视觉
猫头虎10 小时前
手动部署开源OpenClaw汉化中文版过程中常见问题排查手册
人工智能·langchain·开源·github·aigc·agi·openclaw
多恩Stone10 小时前
【3D AICG 系列-9】Trellis2 推理流程图超详细介绍
人工智能·python·算法·3d·aigc·流程图
整得咔咔响10 小时前
贝尔曼最优公式(BOE)
人工智能·算法·机器学习
2501_9469614710 小时前
极简大气创业融资 PPT 模板,适合路演、项目宣讲
人工智能·排序算法
得一录10 小时前
AI 语音助手:如何用大模型优化智能语音交互?
人工智能
玄同76510 小时前
Python 自动发送邮件实战:用 QQ/163 邮箱发送大模型生成的内容
开发语言·人工智能·python·深度学习·机器学习·邮件·邮箱
逸俊晨晖10 小时前
NVIDIA 4090的8路1080p实时YOLOv8目标检测
人工智能·yolo·目标检测·nvidia
Olamyh10 小时前
【手搓 ReAct Agent:告别框架,回归本质】
人工智能·python
seaside200310 小时前
大模型计算量、显存计算方法推导
人工智能·大模型