PyTorch保存网络结构以及参数【 torch.save()、torch.load() 】

2024-01-14 07:33| 来源: 网络整理| 查看: 265

保存网络结构以及参数一、保存方式二、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文件，效果都是一样的。

【本文地址】

公司简介

联系我们