[深度学习] Pytorch(三)—— 多/单GPU、CPU,训练保存、加载模型参数问题
[深度学习] Pytorch(三)—— 多/单GPU、CPU,训练保存、加载预测模型问题
上一篇实践学习中,遇到了在多/单个GPU、GPU与CPU的不同环境下训练保存、加载使用使用模型的问题,如果保存、加载的上述三类环境不同,加载时会出错。就去研究了一下,做了实验,得出以下结论:
多/单GPU训练保存模型参数、CPU加载使用模型
#保存 PATH = 'cifar_net.pth' torch.save(net.module.state_dict(), PATH) #加载 net = Net() net.load_state_dict(torch.load(PATH))
多GPU训练模型、单GPU加载使用模型
#保存 PATH = 'cifar_net.pth' torch.save(net.state_dict(), PATH) #加载 net = Net() net = nn.DataParallel(net) #保存多GPU的,在加载时需要把网络也转成DataParallel的 net.to(device) #放到GPU上 net.load_state_dict(torch.load(PATH)) # 然后测试数据也需要放到GPU上 images, labels = images.to(device), labels.to(device)
多GPU训练保存模型参数、多GPU加载使用模型
#保存 PATH = 'cifar_net.pth' torch.save(net.state_dict(), PATH) #加载 net = Net() net = nn.DataParallel(net) #保存多GPU的,在加载时需要把网络也转成DataParallel的 net.to(device) #放到GPU上 net.load_state_dict(torch.load(PATH)) # 然后测试数据也需要放到GPU上 images, labels = images.to(device), labels.to(device)
可以看到,单GPU和多GPU加载数据的方法其实是一样的,经运行验证,只要按上述代码写,有多个GPU就调用多个,只有一个就调用一个。
另外,保存、加载网络模型还有三种不同的做法
1.保存整个网络模型
2.只保存模型参数(我们用的就是这种)
3.自定义保存
详细方法,请参考:https://blog.csdn.net/Code_Mart/article/details/88254444
相关推荐
深圳克林斯曼 2020-11-09
86377811 2020-11-06
Niteowl 2020-11-05
ohbxiaoxin 2020-11-04
bigquant 2020-11-03
PandaJiong 2020-11-02
fengzhimohan 2020-11-02
xjtukuixing 2020-10-27
深度学习大讲堂 2020-10-14
Trent 2020-10-13
qijiqiguai 2020-10-09
mpkno 2020-09-26
ustbfym 2020-09-21
sunxinyu 2020-09-17
walkpower 2020-09-16
davidsmith 2020-09-04
Icevivian 2020-08-25
QFYJTL 2020-08-23
ohbxiaoxin 2020-08-20