对于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文件,效果都是一样的。

Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐