ONNX
更多代码可以参考我的github：ONNX
ONNX底层实现
ONNX 的存储格式
ONNX 在底层是用 Protobuf 定义的。Protobuf，全称 Protocol Buffer，是 Google 提出的一套表示和序列化数据的机制。使用 Protobuf 时，用户需要先写一份数据定义文件，再根据这份定义文件把数据存储进一份二进制文件。可以说，数据定义文件就是数据类，二进制文件就是数据类的实例。
这里给出一个 Protobuf 数据定义文件的例子：
message Person {
required string name = 1;
required int32 id = 2;
optional string email = 3;
}
这段定义表示在Person这种数据类型中，必须包含 name、id 这两个字段，选择性包含 email字段。根据这份定义文件，用户就可以选择一种编程语言，定义一个含有成员变量name、id、email的Person 类，把这个类的某个实例用 Protobuf 存储成二进制文件；反之，用户也可以用二进制文件和对应的数据定义文件，读取出一个 Person类的实例。
而对于 ONNX ，Protobuf 的数据定义文件在其开源（https://github.com/onnx/onnx/tree/main/onnx），这些文件定义了神经网络中模型、节点、张量的数据类型规范；而二进制文件就是我们熟悉的“.onnx"文件，每一个 onnx 文件按照数据定义规范，存储了一个神经网络的所有相关数据。直接用 Protobuf 生成 ONNX 模型还是比较麻烦的。幸运的是，ONNX 提供了很多实用 API，我们可以在完全不了解 Protobuf 的前提下，构造和读取 ONNX 模型。
ONNX的结构定义
神经网络本质上是一个计算图，计算图的节点是算子，边是参与运算的张量，通过netron可视化ONNX模型，可以直观看到ONNX记录的所有节点的属性信息，并把参与运算的张量信息存储在算子节点的输入输出信息中。ONNX 模型的结构可以用类图大致表示如下：
构造ONNX模型
从上图可以ONNX模型是由以下几个部分构造的：
- ModelProto
- NodeProto
- ValueInfoProto
- GraphProto
尝试完全用 ONNX 的 Python API 构造一个描述线性函数output=a*x+b的ONNX模型，代码如下所示：
import onnx
from onnx import helper
from onnx import TensorProto
# input and output
a = helper.make_tensor_value_info('a', TensorProto.FLOAT, [10, 10])
x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10, 10])
b = helper.make_tensor_value_info('b', TensorProto.FLOAT, [10, 10])
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [10, 10])
# Mul
mul = helper.make_node('Mul', ['a', 'x'], ['c'])
# Add
add = helper.make_node('Add', ['c', 'b'], ['output'])
# graph and model
graph = helper.make_graph([mul, add], 'linear_func', [a, x, b], [output])
model = helper.make_model(graph)
# save model
onnx.checker.check_model(model)
print(model)
onnx.save(model, 'linear_func.onnx')
运行验证ONNX模型：
import onnxruntime
import numpy as np
sess = onnxruntime.InferenceSession('linear_func.onnx')
a = np.random.rand(10, 10).astype(np.float32)
b = np.random.rand(10, 10).astype(np.float32)
x = np.random.rand(10, 10).astype(np.float32)
output = sess.run(['output'], {'a': a, 'b': b, 'x': x})[0]
assert np.allclose(output, a * x + b)
相关案例
pytorch
pytorch转ONNX的API主要讲一下dynamic_axes这个参数。由于Pytorch的Input的flexible特性（见https://stackoverflow.com/questions/66488807/pytorch-model-input-shape），不像Tensorflow那样需要指定input的shape，因此在onnx转换的时候可能需要指定一些动态的输入。如cityscapes的分割模型PIDNet，在尝试转换成ONNX模型的时候，它输入图像的长宽可以是动态变长的，因此可以进行如下设定，表示图像长宽为可变输入。
onnx_file_name = './model/PIDNet_S_Cityscapes_val_13.onnx'
x = torch.randn((1, 3, 12, 12), requires_grad=True)
torch.onnx.export(model,x,onnx_file_name,
export_params=True,
verbose=True,
opset_version=11,
do_constant_folding=True,
dynamic_axes={'input.1':[2,3]})
在netron上查看模型，可以看到图像的长宽为动态输入：
关于pytorch flexibility说明如下：
torch.onnx.export(model, args, f, export_params, verbose=False, training= False)
#Function to Convert to ONNX
def Convert_ONNX():
# set the model to inference mode
model.eval()
# Let's create a dummy input tensor
dummy_input = torch.randn(1, 3, 32, 32, requires_grad=True)
# Export the model
torch.onnx.export(model, # model being run
dummy_input, # model input (or a tuple for multiple inputs)
"ImageClassifier.onnx", # where to save the model
export_params=True, # store the trained parameter weights inside the model file
opset_version=10, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names = ['modelInput'], # the model's input names 可选参数
output_names = ['modelOutput'], # the model's output names 可选参数
dynamic_axes={'modelInput' : {0 : 'batch_size'}, # variable length axes
'modelOutput' : {0 : 'batch_size'}})
print(" ")
print('Model has been converted to ONNX')
以YOLOV4为例。
首先看一下pytorch Yolov4的网络结构,如下代码所示。它的输入是[batchsize, channel, img_height, img_width]，输出是[box, confs]
def transform_to_onnx(weight_file, batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W):
model = Yolov4(n_classes=n_classes, inference=True)
pretrained_dict = torch.load(weight_file, map_location=torch.device('cpu'))
model.load_state_dict(pretrained_dict)
input_names = ["input"]
output_names = ['boxes', 'confs']
dynamic = False
if batch_size <= 0:
dynamic = True
if dynamic:
x = torch.randn((1, 3, IN_IMAGE_H, IN_IMAGE_W), requires_grad=True)
onnx_file_name = "./models/yolov4_-1_3_{}_{}_dynamic.onnx".format(IN_IMAGE_H, IN_IMAGE_W)
dynamic_axes = {"input": {0: "batch_size"}, "boxes": {0: "batch_size"}, "confs": {0: "batch_size"}}
# Export the model
print('Export the onnx model ...')
torch.onnx.export(model,
x,
onnx_file_name,
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=input_names, output_names=output_names,
dynamic_axes=dynamic_axes)
print('Onnx model exporting done')
return onnx_file_name
else:
x = torch.randn((batch_size, 3, IN_IMAGE_H, IN_IMAGE_W), requires_grad=True)
onnx_file_name = "./models/yolov4_{}_3_{}_{}_static.onnx".format(batch_size, IN_IMAGE_H, IN_IMAGE_W)
# Export the model
print('Export the onnx model ...')
torch.onnx.export(model, #pytorch模型
x, #input，目的是onnx用该数据运行一次模型，了解计算过程，进行编译
onnx_file_name, #保存文件名
export_params=True, # 如果为True则导出模型和模型权重,否则导出为训练的模型
opset_version=11, #默认情况下，将模型导出到ONNX子模板的opset版本。默认为9。
do_constant_folding=True, #如果为 True，则在导出期间将恒定折叠优化应用于模型。 常量折叠优化将用预先计算的常量节点替换一些具有所有常量输入的操作。
input_names=input_names, #(字符串列表 ， 默认空列表）–依次分配给图形输入节点的名称 可选参数
output_names=output_names,# (字符串列表 ， 默认空列表）–依次分配给图形输出节点的名称 可选参数
dynamic_axes=None)
print('Onnx model exporting done')
return onnx_file_name
keras
只需要通过pip install keras2onnx就可以安装keras2onnx转换工具，该代码不完整，完整代码请参考：unet
import colorsys
import copy
import time
import numpy as np
from PIL import Image
from nets.unet import Unet as unet
import onnx
import keras2onnx
import onnxruntime
import cv2
class Unet(object):
def __init__(self, **kwargs):
_defaults = {
"model_path" : kwargs['model'],
"model_image_size" : kwargs['model_image_size'],
"num_classes" : kwargs['num_classes']
}
self.onnx_model = kwargs['onnxmodel']
self.__dict__.update(_defaults)
self.generate()
def generate(self):
self.model = unet(self.model_image_size, self.num_classes)
self.model.load_weights(self.model_path)
print('{} model loaded.'.format(self.model_path))
# 导出为onnx模型
onnx_model = keras2onnx.convert_keras(self.model, self.model.name)
onnx.save_model(onnx_model, self.onnx_model)
def detect_image(self, image):
image = image.convert('RGB')
old_img = copy.deepcopy(image)
orininal_h = np.array(image).shape[0]
orininal_w = np.array(image).shape[1]
img, nw, nh = self.letterbox_image(image,(self.model_image_size[1],self.model_image_size[0]))
img = np.asarray([np.array(img).astype(np.float32) /255])
# 使用onnx进行推理预测
sess = onnxruntime.InferenceSession(self.onnx_model)
x = img if isinstance(img, list) else [img]
feed = dict([(input.name, x[n]) for n, input in enumerate(sess.get_inputs())])
pred_onnx = sess.run(None, feed)[0]
pred_onnx = np.squeeze(pred_onnx)
# 使用模型进行推理预测
pr = self.model.predict(img)[0]
pr = pr.argmax(axis=-1).reshape([self.model_image_size[0],self.model_image_size[1]])
pr = pr[int((self.model_image_size[0]-nh)//2):int((self.model_image_size[0]-nh)//2+nh), int((self.model_image_size[1]-nw)//2):int((self.model_image_size[1]-nw)//2+nw)]
seg_img = np.zeros((np.shape(pr)[0],np.shape(pr)[1],3))
for c in range(self.num_classes):
seg_img[:,:,0] += ((pr[:,: ] == c )*( self.colors[c][0] )).astype('uint8')
seg_img[:,:,1] += ((pr[:,: ] == c )*( self.colors[c][1] )).astype('uint8')
seg_img[:,:,2] += ((pr[:,: ] == c )*( self.colors[c][2] )).astype('uint8')
image = Image.fromarray(np.uint8(seg_img)).resize((orininal_w,orininal_h), Image.NEAREST)
blend_image = Image.blend(old_img,image,0.3)
return image, blend_image
推理过程：
# 使用onnx进行推理预测
sess = onnxruntime.InferenceSession(self.onnx_model)
x = img if isinstance(img, list) else [img]
feed = dict([(input.name, x[n]) for n, input in enumerate(sess.get_inputs())])
pred_onnx = sess.run(None, feed)[0]
pred_onnx = np.squeeze(pred_onnx)
PaddleOCR
PaddleOCR 以及PaddleONNX的repo地址：https://github.com/PaddlePaddle/PaddleOCR和https://github.com/PaddlePaddle/Paddle2ONNX。在这里主要是介绍如何将ch_ppocr_server_v2.0_xx的rec模型转换成ONNX模型。介绍这个原因是paddleocr的repo上ch_ppocr_server_v2.0_rec inference模型不支持动态输入，因此需要重新将pretrain model加载导出成inference模型，然后通过paddle2onnx转换成onnx模型。
-
首先需要配置好paddle，paddleocr，paddle2onnx环境。然后将paddleocr的repo git clone到本地。我的目录结构如下图所示。ch_ppocr_server_v2.0_rec 的pretrain 保存在paddleocrmodel下，导出的inference保存在tmp下。
-
执行以下命令：
python tools/export_model.py -c ./configs/rec/ch_ppocr_v2.0/rec_chinese_common_train_v2.0.yml -o \
Global.pretrained_model=./paddleocrmodel/ch_ppocr_server_v2.0_rec_pre/best_accuracy Global.save_inference_dir=./tmp/
- 使用paddle2onnx转换
paddle2onnx --model_dir ./tmp --model_filename inference.pdmodel --params_filename inference.pdiparams --save_file model.onnx --opset_version 10 --enable_onnx_checker True
- 最终使用的推理结果：
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
