Pytorch 保存模型与加载模型

我们经常会看到后缀名为.pt, .pth, .pkl的pytorch模型文件,这几种模型文件在格式上有什么区别吗?
其实它们并不是在格式上有区别,只是后缀不同而已(仅此而已),在用torch.save()函数保存模型文件时,各人有不同的喜好,有些人喜欢用.pt后缀,有些人喜欢用.pth或.pkl.用相同的torch.save()语句保存出来的模型文件没有什么不同。

import torch
#-----方法1
#保存整个网络
torch.save(model, PATH + name) 
# 加载模型
model_dict=torch.load(PATH)

#-----方法2
# 保存网络中的参数, 速度快,占空间少
torch.save(model.state_dict(),PATH + name)
# 加载模型
# 这种方法需要先加载、定义之前的模型
model_dict=model.load_state_dict(torch.load(PATH + name))

参考:
https://zhuanlan.zhihu.com/p/38056115
https://blog.csdn.net/weixin_38145317/article/details/103582549
https://blog.csdn.net/public669/article/details/97752226

你可能感兴趣的:(PyTorch)