上一篇实践学习中,遇到了在多/单个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