使用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网络,它与普通网络有一点细微的不同。废话不多说,直接上代码:
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。