pytorch单机多GPU数据并行训练,torch.DataParallel如何设置随机种子

使用pytorch的单机多GPU数据并行方式训练,torch.DataParallel,如何设置随机种子,使得每次结果可复现?
以下是设置随机种子的代码(两种都试过且无效),我在初始化模型之前调用,但每次结果都不同,模型参数初始化是相同的,不知道哪里设置的不对。请教我这个代码需要怎么调整,或者在程序什么地方设置随机种子?
谢谢!

def set_seed(seed_num=2):
    random.seed(seed_num)
    torch.manual_seed(seed_num)
    np.random.seed(seed_num)
def set_seed(seed_num=2):
    random.seed(seed_num)
    torch.manual_seed(seed_num)
    np.random.seed(seed_num)
    os.environ['PYTHONHASHSEED'] = str(seed_num)
    torch.cuda.manual_seed(seed_num)
    torch.cuda.manual_seed_all(seed_num)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.enabled = True

单GPU环境下随机种子能设置成功

  • 文章:pytorch多gpu DataParallel 及梯度累加解决显存不平衡和显存不足问题 中也许有你想要的答案,请看下吧
  • 除此之外, 这篇博客: Pytorch多GPU计算之torch.nn.DataParallel()中的 3. 如何保存和加载多GPU网络? 部分也许能够解决你的问题, 你可以仔细阅读以下内容或跳转源博客中阅读:
  • Pytorch模型保存和加载的方法,可以看我的这篇博客。在这里不做详细的介绍了,这里只展示如何来保存和加载多GPU网络,它与普通网络有一点细微的不同。废话不多说,直接上代码:

    net = torch.nn.Linear(10,1)  # 先构造一个网络
    net = torch.nn.DataParallel(net, device_ids=[0,3])  #包裹起来
    torch.save(net.module.state_dict(), './networks/multiGPU.pt') #保存网络
    
    # 加载网络
    new_net = torch.nn.Linear(10,1)
    new_net.load_state_dict(torch.load("./networks/multiGPU.pt"))
    

    因为DataParallel实际上是一个nn.Module,所以我们在保存时需要多调用了一个net.module。