系列文章目录

文章目录

系列文章目录
前言
一、ONNX是什么？
二、如何将pth转为onnx
三、ONNX Runtime
- 1. 安装ONNX Runtime包
- 2. 验证模型是否正确
总结

前言

上一节我们已经成功搭建好了模型，并进行训练和保存为pth文件，本小节我们介绍如何将pth文件转为onnx通用格式，以及一些注意事项。

一、ONNX是什么？

ONNX是一种针对机器学习所设计的开放式的文件格式，用于存储训练好的模型。它使得不同的人工智能框架（如Pytorch，MXNet）可以采用相同格式存储模型数据并交互。ONNX的规范及代码主要由微软、亚马逊、facebook和IMB等公司共同开发，以开放源代码的方式托管在Github上。

二、如何将pth转为onnx

在上一节中，我们使用Pytorch框架训练模型并将模型参数保存为pth文件，现在我们要做的是将此转换为ONNX这种统一的格式进行存储。ONNX文件不仅存储了神经网络模型的权重，同时也存储了模型的结构信息以及网络中每一层的输入输出和一些其他的辅助信息。
在Pytorch中已经带有函数torch.onnx.export()可以直接将生成的pth文件转为ONNX文件。

代码如下（示例）：

from model import Net
import torch
path="./model.pth"
onnx_path="./model.onnx"
torch.set_default_tensor_type('torch.FloatTensor')
torch.set_default_tensor_type('torch.cuda.FloatTensor')
net=Net()
net.load_state_dict(torch.load(path, map_location='cpu'))
net.eval()
test_arr = torch.randn(1,1,28,28)
input_names = ['input']
output_names = ['output']
torch.onnx.export(
    net,
    test_arr,
    onnx_path,
    verbose=False,
    opset_version=11,
    input_names=input_names,
    output_names=output_names,
    # dynamic_axes={"input":{3:"width"}}            #动态推理W纬度，若需其他动态纬度可以自行修改，不需要动态推理的话可以注释这行
)
print('->>模型转换成功！')

1.代码解析

1.加载参数

pth_path="./model.pth"
net = Net()
net.load_state_dict(torch.load(pth_path, map_location='cpu'))

其中的map_location参数指的是当加载参数时指定如何重新映射存储位置，可以选择’CPU’和’GPU’。

net.eval()  # 进入推理模式

2.导出ONNX文件

torch.onnx.export(
    net,
    test_arr,
    onnx_path,
    verbose=False,
    opset_version=11,
    input_names=input_names,
    output_names=output_names,
    # dynamic_axes={"input":{3:"width"}}            #动态推理W纬度，若需其他动态纬度可以自行修改，不需要动态推理的话可以注释这行
)

针对torch.onnx.export()函数进行特别声明。

torch.onnx.export(model, # 加载参数后的模型
				  args,  # 给定一组输入，再实际执行一遍模型，即把这组输入对应的计算图记录下来，保存为ONNX格式
				  f, 
				  export_params=True,  # 设置为False不会导出训练好的参数，只会导出一个未训练的模型
				  verbose=False,  # 设置为True将打印导出文件的信息
				  training=<TrainingMode.EVAL: 0>, 
				  input_names=None,  # 计算图中输入节点的名称
				  output_names=None,  # 计算图中输出节点的名称
				  operator_export_type=<OperatorExportTypes.ONNX: 0>, 
				  opset_version=None,  # 算子的版本，不同版本的pytorch默认值不一样
				  do_constant_folding=True, 
			      dynamic_axes=None, 
				  keep_initializers_as_inputs=None, 
				  custom_opsets=None, 
				  export_modules_as_functions=False)

针对dynamic_axes参数进行特别说明，我们首先使用Netron打开为设定dynamic_axes的计算图如下：
在这里插入图片描述
可以看到输入是固定的1(batch_size)x1(channel)x28(height)x28(width)。
我们设置dynamic_axes={“input”:{3:“width”}}，此时我们同样导出计算图，查看此时的计算图如下：

此时，输入的第三个维度不再是28，而是一个动态的数字。总之，如果在推理时你的输入的大小是动态变化的，那么在这里你就需要设置dynamic_axes参数。这里为了方便解释，我们设定输入为固定大小。

2.检验ONNX模型是否正确

代码如下：

import onnx

model_path = './model.onnx'
onnx_model = onnx.load(model_path)
try:
    onnx.checker.check_model(onnx_model)
except onnx.checker.ValidationError as e:
    print("The model is invalid: %s"%e)

else:
    print("The model is valid!")

导入onnx包，如果报错的话需要pip install进行安装。onnx.load函数用于读取一个ONNX模型，onnx.checker.check_model()用于检查模型格式是否正确，如果有错误的话该函数会直接报错。

三、ONNX Runtime

ONNX Runtime是微软开发的一个跨平台机器学习推理加速器，也就是"推理引擎"（同样的还有TensorRT等）。ONNX Runtime是直接对接ONNX的，即ONNX Runtime可以直接读取并运行.onnx文件，而不再需要再把.onnx格式的文件转换成其他格式的文件。如果我们是Pytorch-ONNX-ONNX Runtime这条部署流水线，只要在目标设备中得到.onnx文件，并在ONNX Runtime运行模型，模型部署就算大功告成了。但是本文是使用TensorRT进行推理，所以我们使用ONNX Runtime是为了验证我们导出的ONNX文件的输出是否正确。 上文中是检验模型格式是否正确，两个不相同。

1. 安装ONNX Runtime包

我们要使用GPU进行推理，所以要安装onnxruntime-gpu。如果之前安装的是onnxruntime，卸载重新安装就好。

pip install onnxruntime-gpu

2. 验证模型是否正确

这里我们的思路是：通过ONNX Runtime读取并运行.onnx文件，输入为mnist数据集的测试图片，然后观察输出是不是对应的类别；通过Pytorch读取并运行.pth文件，输入为mnist数据集的测试图片，然后观察输出是不是对应的类别；然后对比两种方式下输出的类别是否相同。整体代码如下：

import cv2
import numpy as np
import onnx
import onnxruntime as ort
import torch
from model import Net
from torchvision import transforms

pth_model_path = './model.pth'
onnx_model_path = './model.onnx'



def pytorch_out(input):
    input = input.cuda()
    model = Net()
    model.load_state_dict(torch.load(pth_model_path))
    model.cuda()
    output = model(input)
    
    output = output.detach().cpu().numpy()
    print(output)
    predict_cla = np.argmax(output)
    print("torch预测是第{}类".format(predict_cla))

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

def onnx_out(input):
    input = to_numpy(input)
    # GPU推理
    sess = ort.InferenceSession('model.onnx', providers=['CUDAExecutionProvider'])
    print("推理方式为：",ort.get_device())
    io_binding = sess.io_binding()

    input_name = sess.get_inputs()[0].name
    output_name = sess.get_outputs()[0].name
    io_binding.bind_cpu_input(input_name,input)
    io_binding.bind_output(output_name)
    sess.run_with_iobinding(io_binding)
    output = io_binding.copy_outputs_to_cpu()[0]
    
    probe = np.squeeze(output[0])  # 输出的是概率
    print(probe)
    predict_cla = np.argmax(probe)  # 概率最大值为对应的类别
    print("onnx预测是第{}类".format(predict_cla))

#读取图片
img = cv2.imread('/home/wjq/pytorch_mnist/mnist_data/mnist_test/7/mnist_test_17.png',0)
# 数据预处理
transform = transforms.Compose([transforms.ToTensor(),
                                            transforms.Normalize((0.1307),(0.3081,))])

input = transform(img)
input_data = torch.unsqueeze(input, dim=-0)  # 加上batch维度 1*1*28*28

pytorch_out(input_data)  # pytorch方法
onnx_out(input_data)   # onnxruntime 方法

先关注文件里面主要的几行代码：

img = cv2.imread('/home/wjq/pytorch_mnist/mnist_data/mnist_test/7/mnist_test_17.png',0)
transform = transforms.Compose([transforms.ToTensor(),
                                            transforms.Normalize((0.1307),(0.3081,))])

input = transform(img)
input_data = torch.unsqueeze(input, dim=-0)  # 加上batch维度 1*1*28*28

pytorch_out(input_data)
onnx_out(input_data)

首先使用cv2.imread读取测试图片，然后经过和训练时相同的数据预处理，又因为是单张图片，所以添加batch_size维度。接写下分别将输入数据通过ONNX Runtime方法和Pytorch方法观察输出。
输出如下图所示：
在这里插入图片描述
可以看到pytorch和onnxruntime的输出相同，以此证明我们使用.pth文件转出的.onnc文件是正确的。

总结

本节我们介绍了如将将.pth文件转为.onnx文件，并验证.onnx文件的格式和内容是否正确，在此之于，我们也介绍了ONNX Runtime这个推理引擎。下一节，我们将介绍如何通过.onnx文件生成TensorRT的引擎文件。

TensorRT推理手写数字分类（二）