Pytorch模型加载函数：torch.load()

梦里梦外; 2023-10-02 10:46 1阅读 0赞

## 一、torch.load()的作用：用来加载torch.save() 保存的模型文件 ##

torch.load()先在CPU上加载，不会依赖于保存模型的设备。如果加载失败，可能是因为没有包含某些设备，比如你在gpu上训练保存的模型，而在cpu上加载，可能会报错，此时，需要使用map\_location来将存储动态重新映射到可选设备上，比如map\_location=torch.device(‘cpu’)，意思是映射到cpu上，在cpu上加载模型，无论你这个模型从哪里训练保存的。

一句话：map\_location适用于修改模型能在gpu上运行还是cpu上运行。

如果map\_location是可调用的，那么对于每个带有两个参数的序列化存储，它将被调用一次:storage和location。存储参数将是存储的初始反序列化，驻留在CPU上。每个序列化存储都有一个与之关联的位置标记，它标识保存它的设备，这个标记是传递给map\_location的第二个参数。内置的位置标签是“cpu”为cpu张量和“cuda:device\_id”(例如:device\_id)。“cuda:2”)表示cuda张力。map\_location应该返回None或一个存储。如果map\_location返回一个存储，它将被用作最终的反序列化对象，已经移动到正确的设备。否则，torch.load()将退回到默认行为，就好像没有指定map\_location一样。

如果map\_location是一个torch.device对象或一个包含设备标签的字符串，它表示所有张量应该被加载的位置。

## 二、使用方式 ##

torch.load(f, map_location=None, pickle_module=<module 'pickle' from '/opt/conda/lib/python3.6/pickle.py'>, **pickle_load_args)

--------------------

参考资料：  
[Pytorch中加载预训练模型/模型参数-torch.load()][Pytorch_-torch.load]  
[torch.load()][torch.load]  
[Pytorch模型加载函数torch.load()][Pytorch_torch.load]  
[pytorch(一)模型加载函数torch.load()][pytorch_torch.load]

[Pytorch_-torch.load]: https://blog.csdn.net/peacefairy/article/details/108977590
[torch.load]: https://blog.csdn.net/weixin_36670529/article/details/106646218
[Pytorch_torch.load]: https://blog.csdn.net/ChaoFeiLi/article/details/116564760
[pytorch_torch.load]: https://blog.csdn.net/pearl8899/article/details/109566084