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模型的转换和简单的使用已经讲完了。祝愿大家共同进步，技术越来越强。