一般使用torch.save()函数将模型保存起来,该函数一般经常与模型的state_dict()方法联合使用。同理,加载模型使用的是torch.load()函数,该函数与模型的load_state_dict()方法联合使用
1.保存模型
torch.save(model.state_dict(), "./model.pth")
执行完该命令后,会在本地目录中生成一个model.pth文件,该文件就是保存好的模型文件
2.载入模型
model.load_state_dict(torch.load("./model.pth"))
使用模型对象的load_state_dict()方法,可以将保存好的模型载入到model里,在执行完该命令之后,model模型中的值将于model.pth文件中的值保持同步
3.将模型载入到指定的硬件设备中
与张量对象类似,模型对象也可以通过cpu()或cuda()方法指定其所在设备。
在使用torch.load()函数时,还可以通过map_location参数指定硬件设备。这样模型就会被载入指定的硬件设备中。例如:
model.load_state_dict(torch.load("./model.pth",map_location={"cuda:1":"cuda:0"}))
这就实现了将模型同时载入了GPU1和GPU0的设备中
但是一般都不是一次性的将模型载入指定硬件设备中,通常都是使用模型的to()方法,将模型复制到指定的设备中去,更利于调试
4.在多卡并行的计算环境中,对模型的保存和载入
在PyTorch的多卡并行计算环境中,模型的内存结构与正常情况下的内存结构是不同的。
如果使用多卡GPU中的单卡训练模型,或是在GPU和CPU设备上切换使用模型,则最好在保存模型时将参数以CPU的方式存储。
#单卡模式
model.cpu().state_dict()
#多卡模式
model.module.cpu().state_dict()
在加载模型上 也可以指定权重到对应的硬件上,例如
将GPU1的权重加载到GPU0上,代码如下:
torch.load("model.pth",map_location={"cuda:1":"cuda:0"})
将GPU的权重加载到CPU上
torch.load("model.pth",map_location = lambda storage, loc:storage)