保存网络结构以及参数
一、保存方式
二、pkl、pth文件区别
2.1 .pkl文件
2.2 .pth文件
对于pytorch保存网络参数,大家一般可以看到有 .pkl文件 以及 .pth文件,对于这两者有什么区别,以及如何保存网络参数等,本文就好好讲述一下。
一、保存方式
首先我们知道不论是保存模型还是参数都需要用到torch.save()。
对于torch.save()有两种保存方式:
只保存神经网络的训练模型的参数,save的对象是model.state_dict();
既保存整个神经网络的的模型结构又保存模型参数,那么save的对象就是整个模型;
Eg. 假设我有一个训练好的模型名叫model,如何来保存参数以及结构?
import torch # 保存模型步骤 torch.save(model, 'net.pth') # 保存整个神经网络的模型结构以及参数 torch.save(model, 'net.pkl') # 同上 torch.save(model.state_dict(), 'net_params.pth') # 只保存模型参数 torch.save(model.state_dict(), 'net_params.pkl') # 同上 # 加载模型步骤 model = torch.load('net.pth') # 加载整个神经网络的模型结构以及参数 model = torch.load('net.pkl') # 同上 model.load_state_dict(torch.load('net_params.pth')) # 仅加载参数 model.load_state_dict(torch.load('net_params.pkl')) # 同上
上面例子也可以看出若使用torch.save()来进行模型参数的保存,那保存文件的后缀其实没有任何影响,.pkl 文件和 .pth 文件一模一样。
二、pkl、pth文件区别
实际上,这两种格式的文件还是有区别的。
2.1 .pkl文件
首先介绍 .pkl 文件,它若直接打开会显示一堆序列化的东西,以二进制形式存储的。如果去 read 这些文件,需要用'rb'而不是'r'模式。
import pickle as pkl file = os.path.join('annot',model.pkl) # 打开pkl文件 with open(file, 'rb') as anno_file: result = pkl.load(anno_file)
或者:
import pickle as pkl file = os.path.join('annot',model.pkl) # 打开pkl文件 anno_file = open(file, 'rb') result = pkl.load(anno_file)
2.2 .pth文件
import torch filename = r'E:\anaconda\model.pth' # 字符串前面加r,表示的意思是禁止字符串转义 model = torch.load(filename) print(model)
但其实不管pkl文件还是pth文件,都是以二进制形式存储的,没有本质上的区别,你用pickle这个库去加载pkl文件或pth文件,效果都是一样的。
model.load_state_dict(state_dict, strict=False)
我们在构造好了一个模型后,可能要加载一些训练好的模型参数。举例子如下:
假设 trained.pth 是一个训练好的网络的模型参数存储
model = Net()是我们刚刚生成的一个新模型,我们希望model将trained.pth中的参数加载加载进来,但是model中多了一些trained.pth中不存在的参数,如果使用下面的命令:
state_dict = torch.load('trained.pth') model.load_state_dict(state_dict)
会报错,说key对应不上,因为model你强人所难,我堂堂trained.pth没有你的那些个零碎玩意,你非要向我索取,我上哪给你弄去。
但是model不干,说既然你不能完全满足我的需要,那么你有什么我就拿什么吧,怎么办呢?下面的指令代码就行了。
model.load_state_dict(state_dict, strict=False)
————————————————
版权声明:本文为CSDN博主「Tchunren」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/t20134297/article/details/110533007
————————————————
版权声明:本文为CSDN博主「Raywit」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/qq_40520596/article/details/106955452