PyTorch编程基础05--保存与载入模型的常用方法

1.保存模型

下面用torch.save()函数将模型保存起来。该函数一般常与模型的state_dict()方法联合使用。

torch.save(model.state_dict(), './model.pth')

执行该命令后，会在本地目录下生成一个model.pth文件。该文件就是保存好的模型文件。

---

2.载入模型

加载模型使用的是torch.load()函数，该函数与模型的load_state_dict()方法联合使用。

model.load_state_dict(torch.load('./model.pth'))
在执行该命令之后，model模型中的值将与model.pth文件中的值保持同步。

---

3.将模型载入指定硬件设备中

与张量对象类似，模型对象也可以通过cpu()或cuda()方法指定其所在设备。
在使用torch.load()函数时，还可以通过map_location参数指定硬件设备。这样模型就会被载入指定的硬件设备中。例如：

model.load_state_dict(torch.load('./model.pth', map_location={'cuda:0':'cuda:1'}))

模型从cuda:0到了cuda:1。
注意：用torch.load()函数一次性将模型载入指定硬件设备中这种操作不是很常用。它可以被拆分成两步：
（1）将模型载入内存
（2）使用模型的to()方法，将模型复制到指定的设备中
这种细粒度的分步载入模型的方法更利于调试

---

4.在多卡并行计算环境中，对模型的保存与载入

在PyTorch的多卡并行计算环境中，模型的内存结构与正常情况下的内存结构是不同的。
如果使用多卡GPU中的单卡进行训练模型，或是在GPU和CPU设备上切换使用模型，则最好在保存模型时将参数以CPU的方式存储

model.cpu().state_dict()  # 单卡模式
model.module.cpu().state_dict()  # 多卡模式

在保存模型的过程中，如果以GPU方式保存模型中的参数，则在模型文件里还会记录该参数所属的GPU序号。
在加载这种模型文件时，系统默认时根据记录中的GPU序号来恢复权重的。如果外界的环境发生了变化（例如在CPU上加载以GPU方法保存的模型文件），则系统无法找到相应的硬件设备，所以就会报错。
当然，在加载模型时也可以指派权重到对应的硬件上。例如：
（1）将GPU1的权重加载到GPU0上

torch.load('./model.pth', map_location={'cuda:1':'cuda:0'})

（2）将GPU的权重加载到CPU上

torch.load('./model.pth', maplocation=lambda storage, loc: storage)