[深度学习] Pytorch—— 多/单GPUCPU，训练保存加载模型参数问题

Posted 2021-03-25 importgpx

[深度学习] 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