本文主要记录onnx的相关知识

onnx介绍

开放神经网络交换 ONNX（Open Neural Network Exchange）是一套表示深度神经网络模型的开放格式，由微软和 Facebook 于 2017 推出。

通俗的说，onnx就是不同模型的中间桥梁。pytorch、tensorflow、mxnet等模型都可以转换为onnx模型，然后onnx模型可以转换为其他模型，例如TensorRT等。

从部署流程上认识onnx

下面是一个常规的神经网络部署流程：

  1. 前期准备
    根据任务问题，我们选择解决的方案，分类、检测、分割等。然后选择合适的模型，例如yolov3、faster rcnn。
  2. 模型训练
    使用pytorch、tensorflow等框架训练模型，训练完成后，保存模型。
    例如训练YOLO模型，我们得到的结果是一个.pt文件
  3. 转换模型
    将训练好的模型转换为onnx模型，方便后续部署。
    onnx可以理解为封装成了一个包，输入输出都是tensor。
  4. 具体部署
    将onnx模型部署到具体的硬件上，例如NVIDIA的tensorrt、高通的SNPE等。

onnx具体部署

下面我们以pytorch模型YOLOv8为例，介绍onnx的部署流程。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 将pytorch的模型转换为onnx模型
import torch

model = YOLO("yolov8s.pt") # 加载模型
output_path = "/PATH/TO/OUTPUT/" # 定义输出路径
input_tensor = np.randn(1, 3, 640, 640) # 定义输入tensor的形状

torch.onnx.export(model,
                input_tensor,
                output_path,
                opset_version=12, # 指定onnx的版本，主要在于一些算子是否支持
                input_names=["images"], # 输入tensor的名称
                output_names=["output"], # 输出tensor的名称
                )
1
2
3
4
5
6
7
8
9
10
11
12
13
# 调用onnxruntime进行推理
import onnxruntime as ort

# onnx inference
onnx_model_path = '/PATH/TO/ONNX_MODEL/'
session = ort.InferenceSession(onnx_model_path)
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 将输入数据转换为onnx网络的输入格式，并且要求是numpy数组
onnx_input = np.random.randn(1, 3, 640, 640) 

result = session.run([output_name], {input_name: onnx_input})

接下来我们具体举例部署mmaction2中的posec3d模型。
mmlab的封装非常严实，因此如果没有官方的部署代码，我们很难直接导出onnx。
首先我们要明确两个东西，输入输出是什么，它的形状是什么的。
这里的输入输出要严格落实到pytorch模型，像mmlab这种进行严格封装的模型，需要进行层层摸索。

例如posec3d的输入是：

1
2
torch.randn(batch, num_joints, num_frames, image_size, image_size)
torch.randn(2, 17, 48, 64, 64)

输出是：

1
2
torch.randn(batch, num_classes)
torch.randn(2, 2)

上述格式就是通过mmlab给的".\tools\deployment\export_onnx_posec3d.py"得到的onnx模型。
要符合onnx的输入输出格式，我们就要将输入输出格式进行转换。

输入的调整直接调用"inference_skeleton"函数，稍微微调即可。输出需要额外调整。

首先是对batch的理解，batch是2num_clips。
就是对输入的一段视频X帧，均匀取num_clips个片段，每个片段长度为clip_length。（这些片段是可以重叠的）
2num_clips是将骨骼信息转换时encoder造成的。

模型会按照batch来进行推理，因此对于输出我们需要进行平均。
github issue