保存与载入模型的常用方法

一般使用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)

本文来自互联网用户投稿，文章观点仅代表作者本人，不代表本站立场，不承担相关法律责任。如若转载，请注明出处。 如若内容造成侵权/违法违规/事实不符，请点击【内容举报】进行投诉反馈！