onnx模型的转换和使用

这篇具有很好参考价值的文章主要介绍了onnx模型的转换和使用。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

目录

前言 

1.torch下将模型转换为onnx模型

2.实际演示转换

3.使用

4.结尾


前言 

ONNX(Open Neural Network Exchange)是一种开放式的文件格式,可以用于保存不同深度学习框架下的网络模型和参数,从而方便将模型进行不同框架下的转换。

1.torch下将模型转换为onnx模型

这里介绍一个函数——torch.onnx.export():

torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, aten=False, export_raw_ir=False, operator_export_type=None, opset_version=None, _retain_param_name=True, do_constant_folding=False, example_outputs=None, strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=None)

其中主要使用的参数含义如下:

model(torch.nn.Module) :需要转换的模型。

args(tuple or torch.Tensor) :模型的输入。因为expoort需要运行模型,所以需要为其提供一个输入。输入的尺寸一致,输入的值随机即可。(注意,除非使用动态轴,否则onnx模型的输入尺寸将被固定为该输入的尺寸)

f:保存路径

export_params(bool):默认为True。如果默认或设置为True,则会将导入模型的权重参数正常保存。如果为False,则会保存一个未训练过的模型。

verbose(bool):如果为True,则打印被导出到标准输出的模型的描述。

input_names(list of str): 按顺序分配名称到中的输入节点。

output_names(list of str):按顺序分配名称到图中的输出节点。

dynamic_axes(dict<string, dict<int, string>> or dict<string, list(int)>, default empty dict):默认情况下,导出的模型将有所有输入和输出张量的形状设置为完全匹配args中给出的形状。要指定张量的轴为动态的(即只有在运行时才知道),要将dynamic_axes设置为一个带schema的字典,如下形式:

torch.onnx.export(SumModule(), (torch.ones(2, 2),), "onnx.pb",
                                  input_names=["x"], output_names=["sum"],
                                  dynamic_axes={
                                      # dict value: manually named axes
                                      "x": {0: "my_custom_axis_name"},
                                      # list value: automatic names
                                      "sum": [0],
                                  })

2.实际演示转换

这里我写了一段深度学习的代码,以此为例来进行实际演示onnx模型的使用。代码如下:

import torch
from torch import nn
from torch.autograd import Variable

class lstm_reg(nn.Module):
    def __init__(self, input_size, hidden_size, output_size=1, num_layers=2):
        super(lstm_reg, self).__init__()

        self.rnn = nn.LSTM(input_size, hidden_size, num_layers,device='cpu')  # rnn
        self.reg = nn.Linear(hidden_size, output_size,device='cpu')  # 回归

    def forward(self, x):
        x, _ = self.rnn(x) 
        s, b, h = x.shape
        x = x.view(s * b, h) 
        x = self.reg(x)
        x = x.view(s, b, -1)
        return x

#随机网络的输入和输出,输入是1000条1x12的向量,输出是1000个值
device=torch.device('cuda')
x=torch.randn(1000,1,12)
x=torch.as_tensor(x,dtype=torch.float64)
y=torch.randn(1000,1,1)
y=torch.as_tensor(y,dtype=torch.float64)

#对这一千条数据划分,七成是训练集,三成是测试集
train_size=int(len(x)*0.7)
test_size=len(x)-train_size
train_x,train_y=x[0:train_size,:,:],y[0:train_size,:,:]
test_x,test_y=x[train_size:,:,:],y[train_size:,:,:]

train=torch.utils.data.TensorDataset(train_x,train_y)
val=torch.utils.data.TensorDataset(test_x,test_y)

net = lstm_reg(12, 100)#net为我们定义的网络,12是输入向量的长度,100是lstm层的深度
net=net.to(device)
criterion = nn.MSELoss()#均方损失函数

optimizer = torch.optim.Adam(net.parameters(), lr=1e-2)#梯度优化方法

if __name__ == '__main__':
    for e in range(1):
        net=net.train()
        train_loader = torch.utils.data.DataLoader(train,
                                                batch_size=32,
                                                num_workers=0,
                                                shuffle=True,
                                                pin_memory=True,
                                                )

        val_loader = torch.utils.data.DataLoader(val,
                                                batch_size=4,
                                                num_workers=0,
                                                pin_memory=True,
                                                )
        for step,(v_x,v_y)in enumerate(train_loader):
            v_x = Variable(v_x).float()
            v_y = Variable(v_y).float()
            v_y=v_y.cuda()
            v_x=v_x.cuda()
    # 前向传播
            out = net(v_x)
            loss = criterion(out, v_y)
    # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print('Epoch: {},step:{} Loss: {:.5f}'.format(e + 1,step,loss.item()))

#本来这里想再写个测试,但是又觉得没有必要...

torch.save(net, "model.pt")

这个代码网络结构为lstm加一层线性层。输入数据和目标输出用了随机的两个张量。

将这段代码运行之后,我们可以得到一个“model.pt”的文件,这是我们在用pytorch进行训练时常常会见到的一种文件格式。下一步我们要做的就是将这个模型转换为onnx文件。代码如下:

import torch

model=torch.load("model.pt",map_location="cpu")#加载需要转换的模型
dummy_input = torch.randn(1,1,12)#随机一个输入
input_names = ["input"]#输入名称
output_names = ["output"]#输出名称
torch.onnx.export(model, dummy_input, r"model.onnx", verbose=False, input_names=input_names, output_names=output_names)

需要注意,在执行torch.save(net, "model.pt")时,保存的是整个网络模型——包括网络框架和权重参数。而大部分时候,为了更灵活对待训练好的模型参数,常常只保存权重参数,使用torch.save(net.state_dict(), 'model.pt')进行保存。这种情况下是不能直接使用上述代码进行模型转换的。

当只保存了网络的权重参数时,我们先要重新加载模型,并导入网络的权重参数。现在我们先将训练代码的最后一行改为,仅保存权重参数:

torch.save(net.state_dict(), "model_only-weights.pt")

随后, 首先从我们定义的lstm_reg中重新加载模型,再进行模型转换,代码如下:

import torch
import demo
from demo import lstm_reg

model=lstm_reg(12,100)#加载模型
weight=torch.load("model_only-weights.pt",map_location="cpu")#加载权重参数
model.load_state_dict(weight)#将权重参数导入模型
dummy_input = torch.randn(1,1,12)
input_names = ["input"]
output_names = ["output"]
torch.onnx.export(model, dummy_input, r"model.onnx", verbose=False, input_names=input_names, output_names=output_names,)

3.使用

这里需要用到一个python库——onnxruntime。直接使用pip install onnxruntime下载即可。使用onnx模型推理代码如下:

import torch
import numpy
import onnxruntime
dummy_input = torch.randn(1,1,12)
session = onnxruntime.InferenceSession("model.onnx")
result = session.run([], {"input":dummy_input.numpy()})
print(result)

输出结果为:

[array([[[-0.09918265]]], dtype=float32)]

4.结尾

至此,onnx模型的转换和简单的使用已经讲完了。祝愿大家共同进步,技术越来越强。 文章来源地址https://www.toymoban.com/news/detail-550372.html

到了这里,关于onnx模型的转换和使用的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包