ONNX 部署 PyTorch 模型

• ONNX RUNTIME 官网主页
• ONNX RUNTIME 官网文档
• 官方托管 Github 仓库: Microsoft/onnxruntime
• ONNX Model Zoo: 预训练的 onnx 模型库
• onnx/models# 目标检测与图像分割 onnx 模型
• (可选）将模型从 PyTorch 导出到 ONNX 并使用 ONNX Runtime 运行
• torch.onnx - 官方文档

Others

• torchvision 官方 Mask RCNN 转 ONNX

• onnxruntime 调用 AI 模型的 python 和 C++ 编程

• 深度学习模型转换与部署那些事(含 ONNX 格式详细分析)

• tenglike1997/onnxruntime-projects - Github

• ONNX 动态输入和动态输出问题

• onnxruntime 的 c++ 使用(含多输入多输出问题)

• MIMO 模型的导入与运行 - Github

ONNX C++ API 说明

1. 对每个输入节点需要创建一个 Ort::Value 类型的变量, 该变量可以是具有动态维度的张量. 例如可以一次性输入 1 张图片, 也可以一次性输入 batch_size=100 张图片.

2. 创建 Ort::Value 类型的变量需要通过 Ort::Value::CreateTensor 函数来创建, 该函数有几种重载形式, 对于如下这种 CreateTensor 形式而言

static Ort::Value Ort::Value::CreateTensor<T>(const OrtMemoryInfo* info, T* p_data, size_t p_data_element_count, const int64_t* shape, size_t shape_len);

创建张量就需要将具体的数据组织成一维数组, 将首地址传入 T* p_data,
数据元素总数传入 size_t p_data_element_count,
张量形状组织成一维数组, 首地址传入 const int64_t* shape,
张量总维数传入size_t shape_len.

例如:
要将 100 张 32*64 的 RGB 图片创建为张量, 那么原始数据形状可能为 [100,3,32,64], 需要将其变成一维的数组std::vector<float> p_data(100*3*32*64), 然后将数组首地址p_data 传入. 那么元素总数就是 p_data_element_count=100*3*32*64=614400. 为了让ONNX RUNTIME 明确数据的原始格式, 需要创建 std::vector<int64_t> shape = {100,3,32,64}, 然后将shape 作为维度数组的首地址传入. 最后该数据原本是 4 维的, 因此shape_len=4.

3. 如果是多输入的, 那么首先创建多个 Ort::Value 类型的变量, 然后创建 std::vector<Ort::Value> ort_inputs, 再通过移动构造将这些变量压入vector 容器中.

ort_inputs.push_back(std::move(Tensor_1));
ort_inputs.push_back(std::move(Tensor_2));

这样, 形式上可以将任意多个 Ort::Value 变量在内存中放在一起了.

4. 在进行模型推理时, 调用 session.Run() 函数. 该函数有 三种重载版本 .
   第一种是只需要给输入变量预分配内存, 将输出变量通过函数返回值返回;
   第二种需要给输入输出变量都预分配内存, 没有返回值;
   第三种将输入输出进行IoBinding&, 没有返回值.

对于第一种版本:

std::vector<Value> Run(const RunOptions& run_options, const char* const* input_names, const Value* input_values, size_t input_count,
                         const char* const* output_names, size_t output_count);

input_names是输入节点名的字符串数组的首地址, 例如 std::vector<const char*> names = {"input_1, "input_2"}. 那么将names 传入即可. input_values是要传入的若干个张量的首地址, 如果只传入一个张量, 那就对这个张量取地址 &Tensor_1 传入即可; 如果要传入多个张量, 那么将这多个张量放在一个 vector 中, 并将该 vector 首地址传入即可 (二者实际是一致的).
input_count 是要输入张量的个数.
output_names,output_count同理.

5. 在对输出变量进行解码时, 首先要按照输出节点数分为多个张量, 视该张量数据类型不同, 以不同数据类型取出其中的数据.
   例如下面代码块中,type就是数据类型.

   // print output node types
   Ort::TypeInfo type_info = session.GetOutputTypeInfo(i);
   auto tensor_info = type_info.GetTensorTypeAndShapeInfo();
   
   ONNXTensorElementDataType type = tensor_info.GetElementType();
   printf("Output %d : type=%d\n", i, type);

其中 ONNXTensorElementDataType 是一个枚举类型, 具体定义如下:

// https://github.com/microsoft/onnxruntime/blob/master/include/onnxruntime/core/session/onnxruntime_c_api.h#L93
// Copied from TensorProto::DataType
// Currently, Ort doesn't support complex64, complex128
typedef enum ONNXTensorElementDataType {
  ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED,
  ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,   // maps to c type float
  ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8,   // maps to c type uint8_t
  ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8,    // maps to c type int8_t
  ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16,  // maps to c type uint16_t
  ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16,   // maps to c type int16_t
  ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32,   // maps to c type int32_t
  ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64,   // maps to c type int64_t
  ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING,  // maps to c++ type std::string
  ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL,
  ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16,
  ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE,      // maps to c type double
  ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32,      // maps to c type uint32_t
  ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64,      // maps to c type uint64_t
  ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64,   // complex with float32 real and imaginary components
  ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128,  // complex with float64 real and imaginary components
  ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16     // Non-IEEE floating-point format based on IEEE754 single-precision
} ONNXTensorElementDataType;

对于不同的数据类型, 对应的内存块中存储的是整型变量或者浮点型变量, 应当进行不同形式的解码, 如下.

float* boxes_ptr = output_tensors[0].GetTensorMutableData<float>();
int* labels_ptr = output_tensors[1].GetTensorMutableData<int>();

解码后得到的是一维数据, 还应当依据输出张量的维度信息进行重新组织, 得到合理的输出.

线性回归例程

这里通过一个简单的线性回归模型来说明 onnx 模型文件的保存与加载

上图所示是一个简单的线性回归模型, 用 python+pytorch 创建并训练, 真实的回归模型为 y1 = 12*x+9, y2 = 4*x+9. 通过大量的随机数进行训练使其能够稳定预测输出值. 然后通过torch.onnx 将模型输出为 model.onnx 文件, 并再次通过 python+onnxruntime 读取该模型以验证模型文件是没有问题的.

然后通过 CPP+onnxruntime_cxx_api 读取该模型文件, 并输入 batch_size=6 个数据, 从得到的 [batch_size=6,2] 个数据来看, 可以较为准确地预测输出值.

标量输入:x, 向量输出 y = [y1,y2]
y1 = k1*x+b1
y2 = k2*x+b2

• 通过 pytorch 创建模型、训练并保存 onnx 模型文件
  • Python 脚本: linear_regression2d.py
  • 运行输出: linear_regression2d_py_output.txt
• 通过 cpp 加载 onnx 模型文件, 并进行推理
  • CPP 程序: linear_regression2d.cpp
  • 运行输出: linear_regression2d_cpp_output.txt

MaskRCNN 例程

当前有一个任务是要将 python+pytorch 编写的 maskrcnn 模型部署到 C++ 环境中, 主要路线是用 pytorch 内置的 torch.onnx 将模型输出为 *.onnx 文件, 然后再用 onnxruntime_cxx_api 导入模型进行推理.

示例图片如下: demo.jpg

1. 调用 pytorch 内置模型 maskrcnn_resnet50_fpn(包括预训练参数). 使用torch.jit.script 直接输出序列化模型文件 model.pt.
   如果用 C++ 版本的 libtorch+torchvision 也许是可以直接导入该模型文件并进行推理的, 但是 torchvision 的编译过程踩坑不断, 已弃疗!

• 创建并输出 pt 模型文件

1. 调用 torch.onnx 输出序列化模型文件 model.onnx, 输出时需要提供一个示例样本, 理论上可以用随机数torch.randn, 这不影响模型导出以及后续的加载. 只是随机生成的样本可能经过模型推理后输出为空, 在python 环境中也许还比较容易发现异样, 但在用 C++ 编程加载模型并进行推理的时候如果也用随机样本, 并且输出为空, 会让人怀疑人生的!

requirements:

> conda install -c conda-forge onnx 

• 创建并输出 ONNX 模型文件

1. 仍旧是用 python 加载模型文件 model.onnx, 并用onnxruntime 进行模型推理. 原始模型推理得到的输出是 batch_size 个字典组成的列表. 每个字典是对应某张图片的四个输出: "boxes","labels","scores","masks". 而 onnxruntime 推理得到的输出似乎直接就全都是列表(待证!).

Requirements:

> pip install onnxruntime

• Python 导入模型前向推理

4. 以上是在 python 中进行模型导出与加载, 接下来要换用 C++ 接口实现.

环境: Visual Studio 2017, Windows 10 x64.

新建 VS 工程 -> 项目 -> 管理 NuGet 程序包 (N)-> 浏览 -> 输入 onnx: 选择 Microsoft.ML.OnnxRuntime+Microsoft.ML.OnnxRuntime.MKLML 进行安装 (每新建一个工程就要安装一次). 如果要用GPU 加速, 则应该下载 GPU 版本, 并且似乎 CPU 与GPU版本不可共存. 并且 Python 环境中也要对应, 使用一个就要卸载另一个.(这个后续再去仔细研究, 暂且仅用 CPU 版本)

参考:

• 如何在 VS 项目中添加 onnx runtime 包

5. 读取图片

Note: 我首先将 demo_maskrcnn.jpg resize 成了 500*667 大小的图片demo_resize.jpg, 且此后所有操作都是针对处理后的图片进行的.

from PIL import Image
img = Image.open('demo.jpg')
img = img.resize((500,667))
img.save('demo_resize.jpg')

• 图片读取与处理 .cpp
• 图片读取与处理 .h

6. ONNX C++ API调用模型进行前向推理

• C++ 模型前向推理

用 python(左) 和cpp(右)推理得到的模型输出结果依次如下图所示, 虽然有所区别, 但如果将置信度设置在 90% 以上, 则二者都检测出了图片中人物的全身. 此外, 这里试验用的模型是 pytorch预训练模型, 并没有针对此类图片进行专门训练. 测试出现偏差也情有可原.

• CPP 模型推理代码封装

  • main.cpp
  • MaskRCNN_Session.cpp
  • MaskRCNN_Session.h

• 本文例程打包下载: https://github.com/siyouluo/Storage/releases/download/v1.0.0/onnx.zip