pytorch固定随机种子-训练稳定复现

import torch

import random

import numpy as np

import os

def set_seed(seed=1): # seed 的数值可以随意设置，本人不清楚有没有推荐数值

    random.seed(seed)

    np.random.seed(seed)

    torch.manual_seed(seed)

    #根据文档，torch.manual_seed (seed) 应该已经为所有设备设置 seed

    #但是 torch.cuda.manual_seed (seed) 在没有 gpu 时也可调用，这样写没什么坏处

    torch.cuda.manual_seed(seed)

    torch.cuda.manual_seed_all(seed)

    os.environ['PYTHONHASHSEED'] = str(seed)

    #cuDNN 在使用 deterministic 模式时（下面两行），可能会造成性能下降（取决于 model）

    torch.backends.cudnn.deterministic = True

    torch.backends.cudnn.benchmark = False

GLOBAL_SEED = 1

GLOBAL_WORKER_ID = None

def worker_init_fn(worker_id):

    global GLOBAL_WORKER_ID

    GLOBAL_WORKER_ID = worker_id

    set_seed(GLOBAL_SEED + worker_id)

train_loader = DataLoaderX(

        local_rank=local_rank, dataset=trainset, batch_size=cfg.batch_size,

        sampler=train_sampler, num_workers=0, pin_memory=True, drop_last=True,

        worker_init_fn=worker_init_fn)

rng = numpy.random.RandomState(1)

rng.uniform(0,1,(2,3))

rng = numpy.random.RandomState(1)

rng.uniform(0,1,(2,3))