上一篇实践学习中,遇到了在多/单个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
本站文章如无特殊说明,均为本站原创,如若转载,请注明出处:[深度学习] Pytorch(三)—— 多/单GPU、CPU,训练保存、加载模型参数问题 - Python技术站