领先的手机网站设计,网络推广十大平台,做网站关键字,三类人员 网站开发1 pytorch保存和加载模型的三种方法
PyTorch提供了三种方式来保存和加载模型#xff0c;在这三种方式中#xff0c;加载模型的代码和保存模型的代码必须相匹配#xff0c;才能保证模型的加载成功。通常情况下#xff0c;使用第一种方式#xff08;保存和加载模型状态字典…1 pytorch保存和加载模型的三种方法
PyTorch提供了三种方式来保存和加载模型在这三种方式中加载模型的代码和保存模型的代码必须相匹配才能保证模型的加载成功。通常情况下使用第一种方式保存和加载模型状态字典更加常见因为它更轻量且不依赖于特定的模型类。
1.1 仅保存和加载模型参数推荐
1.1.1 保存模型参数
import torch
import torch.nn as nnmodel nn.Sequential(nn.Linear(128, 16), nn.ReLU(), nn.Linear(16, 1))# 保存整个模型
torch.save(model.state_dict(), sample_model.pt)
1.1.2 加载模型参数
import torch
import torch.nn as nn# 下载模型参数 并放到模型中
loaded_model nn.Sequential(nn.Linear(128, 16), nn.ReLU(), nn.Linear(16, 1))
loaded_model.load_state_dict(torch.load(sample_model.pt))
print(loaded_model)
显示如下
Sequential((0): Linear(in_features128, out_features16, biasTrue)(1): ReLU()(2): Linear(in_features16, out_features1, biasTrue)
)
net.state_dict()在PyTorch中Module 的可学习参数 (即权重和偏差)模块模型包含在参数中 (通过 model.parameters() 访问)。state_dict 是一个从参数名称隐射到参数 Tesnor 的有序字典对象。只有具有可学习参数的层(卷积层、线性层等) 才有 state_dict 中的条目。
1.2 保存和加载整个模型
1.2.1 保存整个模型
import torch
import torch.nn as nnnet nn.Sequential(nn.Linear(128, 16), nn.ReLU(), nn.Linear(16, 1))# 保存整个模型包含模型结构和参数
torch.save(net, sample_model.pt)
1.2.2 加载整个模型
import torch
import torch.nn as nn# 加载整个模型包含模型结构和参数
loaded_model torch.load(sample_model.pt)
print(loaded_model)
显示如下
Sequential((0): Linear(in_features128, out_features16, biasTrue)(1): ReLU()(2): Linear(in_features16, out_features1, biasTrue)
)
1.3 导出和加载ONNX格式模型
1.3.1 保存模型
import torch
import torch.nn as nnmodel nn.Sequential(nn.Linear(128, 16), nn.ReLU(), nn.Linear(16, 1))input_sample torch.randn(16, 128) # 提供一个输入样本作为示例
torch.onnx.export(model, input_sample, sample_model.onnx)
1.3.2 加载模型
import torch
import torch.nn as nn
import onnx
import onnxruntimeloaded_model onnx.load(sample_model.onnx)
session onnxruntime.InferenceSession(sample_model.onnx)
print(session)
2 模型保存与加载使用的函数
2.1 保存模型函数torch.save
将对象序列化保存到磁盘中该方法原理是基于python中的pickle来序列化各种Models,tensors,dictionaries 都可以使用该方法保存。保存的模型文件名可以是.pth, .pt, .pkl。
def save(obj: object,f: FILE_LIKE,pickle_module: Any pickle,pickle_protocol: int DEFAULT_PROTOCOL,_use_new_zipfile_serialization: bool True
) - None:
obj保存的对象f一个类似文件的对象必须实现写入和刷新或字符串或操作系统。包含文件名的类似路径对象pickle_module用于挑选元数据和对象的模块pickle_protocol可以指定以覆盖默认协议
2.2 加载模型函数torch.load
def load(f: FILE_LIKE,map_location: MAP_LOCATION None,pickle_module: Any None,*,weights_only: bool False,**pickle_load_args: Any
) - Any:
f类文件对象 (返回文件描述符)或一个保存文件名的字符串map_location一个函数或字典规定如何映射存储设备torch.device对象pickle_module用于 unpickling 元数据和对象的模块 (必须匹配序列化文件时的 pickle_module )
2.3 加载模型参数torch.nn.Module.load_state_dict
序列化 (Serialization)是将对象的状态信息转换为可以存储或传输的形式的过程。 在序列化期间对象将其当前状态写入到临时或持久性存储区。以后可以通过从存储区中读取或反序列化对象的状态重新创建该对象。
def load_state_dict(self, state_dict: OrderedDict[str, Tensor],strict: bool True):
state_dict保存 parameters 和 persistent buffers 的字典strict可选bool型。state_dict 中的 key 是否和 model.state_dict() 返回的 key 一致。
2.4 状态字典state_dict
函数作用是“获取优化器当前状态信息字典”在神经网络中模型上训练出来的模型参数也就是权重和偏置值。在Pytorch中定义网络模型是通过继承torch.nn.Module来实现的。其网络模型中包含可学习的参数weights, bias, 和一些登记的缓存如batchnorm’s running_mean 等。模型内部的可学习参数可通过两种方式进行调用
通过model.parameters()这个生成器来访问所有参数。通过model.state_dict()来为每一层和它的参数建立一个映射关系并存储在字典中其键值由每个网络层和其对应的参数张量构成。
def state_dict(self, destinationNone, prefix, keep_varsFalse):
除模型外优化器对象(torch.optim)同样也有一个状态字典包含的优化器状态信息以及使用的超参数。由于状态字典属于Python 字典因此对 PyTorch 模型和优化器的保存、更新、替换、恢复等操作都比较便捷。
3 指定map_location加载模型
采用仅加载模型参数的方式指定设备类型进行模型加载代码如下
model_path /opt/sample_model.pthdevice torch.device(cuda if torch.cuda.is_available() else cpu)
map_location torch.device(device)model.load_state_dict(torch.load(self.model_path, map_locationself.map_location))
