ONNX知识点

本文最后更新于 2025年9月3日 下午

本文主要分享了 ONNX 的基础知识。

概述

ONNX(Open Neural Network Exchange)是一个开放的神经网络交换格式,用于在不同的深度学习框架之间进行模型互操作。它由微软和 Facebook 在 2017 年联合推出,旨在让开发者能够更方便地在不同的框架(如 PyTorch、TensorFlow、MXNet、Caffe2 等)之间转换模型。

特点

  1. 开放标准
    ONNX 是一个开放标准,允许开发者定义神经网络模型的表示形式,使用标准化的格式存储和共享模型。
  2. 跨框架互通
    将不同框架(如 PyTorch、TensorFlow、Caffe2、MXNet 等)训练出的模型,转为统一的 ONNX 格式文件(.onnx),方便在不同环境下直接部署或使用,避免被特定框架绑定。
  3. 可扩展性
    ONNX 支持多种操作(operators)和扩展,允许开发者根据需求创建自定义操作。
  4. 生态丰富
    多种工具和运行时(如 ONNX Runtime、TensorRT、OpenVINO 等)支持 ONNX 模型,加速推理。

主要作用

ONNX(Open Neural Network Exchange) 是一种开放的神经网络模型交换格式。它的主要用途是作为不同深度学习框架之间的“中间桥梁”,让模型可以在不同平台、框架和硬件之间迁移和部署,大大提升了模型的可移植性和工程效率。

ONNX 并不是直接本身“支持”比 PyTorch 或 TensorFlow 更多的平台和硬件,而是因为它作为通用模型格式,能够被更多第三方推理引擎和硬件生态支持。因此,使用 ONNX 能极大提升模型在不同平台和硬件上的迁移能力和部署灵活性,让模型脱离原始训练框架的限制,易于集成和工程落地。

  • 模型互操作:支持 PyTorch、TensorFlow、Keras、Caffe2、MXNet 等多种主流深度学习框架之间模型的互相转换和迁移。
  • 统一推理接口:通过 ONNX Runtime、TensorRT(TensorRT 从 5.x 版本开始原生支持 ONNX 格式。你可以直接用 TensorRT 的 API 加载、优化并在 NVIDIA GPU 上加速 ONNX 模型推理) 等推理引擎加速和统一模型部署,支持多种硬件(CPU、GPU、FPGA、AI 芯片等)。
  • 简化生产部署流程:模型训练和推理可以分离,方便团队协作和跨平台部署。

平台主要指 操作系统环境应用场景,包括但不限于:

  • 操作系统:Windows、Linux、macOS、Android、iOS
  • 云平台:AWS、Azure、Google Cloud、Alibaba Cloud 等
  • 边缘设备:嵌入式系统、移动设备、IoT 设备
  • Web 环境:浏览器(通过 ONNX.js 或 WebAssembly)

框架指的是 深度学习/机器学习工具和库,ONNX 支持主流框架之间的模型转换和互操作,例如:

  • PyTorch:支持模型导出为 ONNX
  • TensorFlow:通过 tf2onnx、onnx-tensorflow 等工具支持转换
  • Keras:间接支持(通过 TensorFlow 或其他工具转换为 ONNX)
  • Caffe2
  • MXNet
  • Scikit-learn:部分模型可转为 ONNX
  • LightGBM、XGBoost:部分机器学习模型支持 ONNX 导出
  • 其他支持 ONNX 的框架:如 PaddlePaddle、CoreML 等(部分支持)

硬件是指 模型实际推理运行的物理设备类型,ONNX 通过 Runtime 支持多种硬件加速:

  • CPU:x86、ARM 等架构
  • GPU:NVIDIA CUDA(通过 TensorRT、ONNX Runtime、DirectML)、AMD GPU
  • FPGA:部分 FPGA 厂商支持 ONNX(如 Xilinx)
  • AI 专用芯片:如华为昇腾、寒武纪、Apple Neural Engine(部分支持 ONNX)
  • 移动芯片:高通 Snapdragon NPU、苹果 A 系列等
  • TPU:谷歌 TPU(间接支持)

典型场景

  • 跨框架协作:研究团队用 PyTorch 训练,生产环境用 TensorFlow Serving 部署,或者反过来,ONNX 可以作为中间格式实现无缝迁移。
  • 硬件加速部署:AI 芯片、嵌入式设备、云端推理等场景下,很多硬件厂商只支持 ONNX 格式,直接利用 ONNX Runtime 或 TensorRT 做高效推理。
  • 模型集成与复用:第三方模型(如 HuggingFace、Open Model Zoo)通常提供 ONNX 格式,便于集成到自己的系统中。
  • 减少重复开发:不用为每种推理后端、硬件写专用转换代码,ONNX 提供统一标准。

常用流程

  1. 训练:在如 PyTorch 中训练模型。
  2. 导出模型:将训练好的模型导出为 ONNX 格式。
  3. 部署:用 ONNX Runtime 等运行时/推理引擎部署于各种硬件/操作系统/云端或边缘设备上。

是否需要重新实现网络结构?

概述

使用 ONNX 的确会涉及模型的网络结构,但具体是否需要重新实现网络结构,取决于以下几个因素:

  1. 是否已有 PyTorch 模型
    如果你已经有了一个完整的 PyTorch 模型(权重+网络结构),并且它可以正常运行,你可以直接将 PyTorch 模型导出为 ONNX 格式,而不需要重新实现网络结构。ONNX 的主要作用是作为一个中间格式,用于存储和交换模型的定义和权重。

    检查模型保存方式

    判断.pt 文件是否同时包含网络结构和权重

  2. ONNX 导出是否支持所有算子
    PyTorch 提供了 torch.onnx.export 方法,可以将模型导出为 ONNX 格式。然而,某些复杂或自定义的 PyTorch 操作(算子)可能不被 ONNX 支持。如果模型中使用了这些不支持的操作,你可能需要对模型进行部分修改,或者在导出后手动重新定义这些部分的网络结构。

  3. ONNX 推理引擎的支持
    使用 ONNX Runtime 或其他推理引擎进行模型部署时,推理引擎需要能够解析并执行 ONNX 模型中的所有操作。如果你的模型中包含不被推理引擎支持的操作,则可能需要重新实现这些部分的网络结构。

检查是否需要重新实现网络结构

  • 导出并验证 ONNX 模型
    在导出模型后,使用工具(如 onnx.checker 或 ONNX Runtime)检查模型的有效性。如果导出过程中或者验证时出现不支持的算子,可能需要修改网络结构。
  • 观察导出日志
    torch.onnx.export 会在导出过程中生成日志,如果某些操作不被支持,日志中会有相关提示。
  • 使用替代操作或自定义算子
    如果导出过程中某些操作不被支持,可以尝试用 ONNX 支持的操作替代,或者在 ONNX 中定义自定义算子。

基本概念

数据类型

种类

onnx 的数据类型,共有 16 种

  • elem_type: 1 --> float32
  • elem_type: 2 --> uint8
  • elem_type: 3 --> int8
  • elem_type: 4 --> uint16
  • elem_type: 5 --> int16
  • elem_type: 6 --> int32
  • elem_type: 7 --> int64
  • elem_type: 8 --> string
  • elem_type: 9 --> boolean
  • elem_type: 10 --> float16
  • elem_type: 11 --> float64
  • elem_type: 12 --> uint32
  • elem_type: 14 --> uint64
  • elem_type: 15 --> complex128
  • elem_type: 16 --> bfloat16

查看

1
2
3
4
5
6
7
8
9
10
11
import onnx
import onnxruntime
import numpy as np

# 加载 ONNX 模型
onnx_model_path = "./weights/superpoint.onnx"
onnx_model = onnx.load(onnx_model_path)

# 检查 ONNX 模型的输出类型
for o in onnx_model.graph.output:
print(o.name, o.type)

需要提前确定吗?

ONNX 模型的输出类型在 C++和 Python 中都需要提前确定,但两者的“敏感度”不同:

  1. C++不像 Python 有动态类型,你必须明确知道输出的数据类型(如 float*int64_t* 等),否则会崩溃或数据错乱。
  2. Python 的 ONNXRuntime 会自动把输出转成 numpy 数组,类型自动匹配,你一般不用手动指定类型。但如果你要做后续处理(比如拼接、转存、可视化),也需要知道输出的 shape 和 dtype,否则容易出错或报错。

PyTorch-> ONXX

转换的好处和优势

将 PyTorch 模型转换为 ONNX(Open Neural Network Exchange)模型有以下好处和优势:

  1. 跨平台部署
    ONNX 是一个开放的深度学习模型交换格式,支持多种框架(如 PyTorch、TensorFlow、Caffe2、MXNet 等)和推理引擎(如 ONNX Runtime、TensorRT、OpenVINO 等)。将 PyTorch 模型转为 ONNX 后,可以在不同的平台、硬件和框架上进行部署,大大提高了模型的通用性和适应性。
  2. 推理加速
    许多高性能推理引擎(如 ONNX Runtime、TensorRT 等)支持对 ONNX 模型进行优化,可以充分利用 CPU、GPU、FPGA、NPU 等硬件资源,实现更快的推理速度和更低的延迟。
  3. 硬件兼容性
    ONNX 模型支持多种硬件后端,方便在嵌入式设备、移动端、云端等不同硬件环境下部署,提升模型的可移植性。
  4. 模型优化和量化
    ONNX 生态提供了丰富的模型优化工具(如 onnxoptimizer、onnxruntime-tools),支持模型剪枝、量化、融合等操作,进一步提升模型推理性能和减少模型体积。
  5. 标准化和可扩展性
    ONNX 作为开放标准,定义了统一的运算符集合和模型格式,使模型的结构清晰、易于解析和维护。同时,ONNX 社区活跃,持续支持新算子和特性。
  6. 便于集成到生产环境
    许多生产级 AI 推理服务和平台(如 Azure、AWS、NVIDIA Triton 等)原生支持 ONNX 格式,便于将 PyTorch 训练好的模型快速集成到实际应用中。
  7. 便于模型测试和验证
    ONNX 模型可以在多个平台上进行一致性测试,确保模型行为在不同环境中的一致性,降低因框架差异导致的 bug 风险。

总结
PyTorch 模型转为 ONNX 后,可以实现跨平台、跨框架、跨硬件的灵活部署,同时获得更好的推理性能和硬件兼容性,是模型落地和生产化的重要步骤。

ONNX 模型可以转换成 PyTorch 或 TensorFlow 模型吗?

  • 通常情况:ONNX 支持从 PyTorch、TensorFlow 等导出到 ONNX,但从 ONNX 转回原框架不一定总是可行。
  • PyTorch → ONNX:PyTorch 官方支持将模型导出为 ONNX 格式(torch.onnx.export())。
  • TensorFlow → ONNX:通过工具如 tf2onnx
  • ONNX → PyTorch/TensorFlow:理论上可用一些第三方工具尝试转换(如 onnx2pytorchonnx-tf),但不保证所有算子都能还原,复杂模型可能不完全兼容,通常只推荐做推理,不推荐做训练。

导出 PyTorch 模型到 ONNX 格式

一般流程

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import torch
from demo_superpoint import SuperPointNet # 确保导入你的网络定义

if __name__ == '__main__':
# 1. 定义网络结构
model = SuperPointNet()

# 2. 加载原始权重文件
weights_path = "./weights/superpoint_v1.pth" # 原始权重文件路径
state_dict = torch.load(weights_path)
model.load_state_dict(state_dict)

# 3. 设置为评估模式
model.eval()

# 4. 导出为 ONNX 模型
onnx_output_path = "./weights/superpoint.onnx" # ONNX 模型输出路径
example_input = torch.randn(1, 1, 240, 320) # 示例输入,假设输入大小为 240x320
torch.onnx.export(
model, # 要导出的模型
example_input, # 示例输入
onnx_output_path, # 输出文件路径
export_params=True, # 保存模型参数
opset_version=11, # ONNX opset 版本
do_constant_folding=True, # 是否执行常量折叠优化
input_names=['input'], # 输入张量的名称
output_names=['semi', 'desc'], # 输出张量的名称
dynamic_axes={ # 动态轴支持
'input': {0: 'batch_size'}, # 输入的第 0 维是动态的(batch size)
'semi': {0: 'batch_size'}, # 输出的第 0 维是动态的(batch size)
'desc': {0: 'batch_size'} # 输出的第 0 维是动态的(batch size)
}
)

print(f"ONNX model exported to {onnx_output_path}")

示例输入的大小(example_input)需要与模型的实际输入大小一致。

支持动态输入

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 导出 ONNX 模型
torch.onnx.export(
model, # 要导出的模型
example_input, # 示例输入
onnx_output_path, # 输出文件路径
export_params=True, # 保存模型参数
opset_version=11, # ONNX opset 版本
do_constant_folding=True, # 是否执行常量折叠优化
input_names=['input'], # 输入张量的名称
output_names=['semi', 'desc'], # 输出张量的名称
dynamic_axes={ # 动态轴支持
'input': {2: 'height', 3: 'width'}, # 输入的第 2 和第 3 维是动态的(高度和宽度)
'semi': {2: 'height', 3: 'width'}, # 输出 semi 的第 2 和第 3 维是动态的
'desc': {2: 'height', 3: 'width'} # 输出 desc 的第 2 和第 3 维是动态的
}
)
  • dynamic_axes 参数用于指定输入和输出张量的哪些维度是动态的。
  • 示例输入的形状([1, 1, 240, 320])只是一个占位符,用于跟踪模型的计算图。
  • 导出的 ONNX 模型将支持动态大小的输入,而不仅仅是固定的 240x320

验证 ONNX 模型

  • 使用网站 Netron 验证导出的 ONNX 模型是否正确。

  • 使用 ONNX Runtime 验证模型:

    1. 安装依赖:

      1
      2
      3
      4
      5
      pip install numpy
      pip install protobuf
      sudo apt-get install protobuf-compiler libprotoc-dev
      export CMAKE_ARGS="-DONNX_USE_PROTOBUF_SHARED_LIBS=ON"
      pip install onnx onnxruntime
    2. 示例代码:

      1
      2
      3
      4
      5
      6
      7
      8
      9
      10
      11
      12
      13
      14
      15
      16
      17
      18
      19
      20
      21
      22
      23
      24
      25
      26
      27
      28
      29
      30
      31
      import onnx
      import onnxruntime
      import numpy as np

      # 加载 ONNX 模型
      onnx_model_path = "./weights/superpoint.onnx"
      onnx_model = onnx.load(onnx_model_path)

      # 检查模型是否有效
      onnx.checker.check_model(onnx_model)
      print("ONNX model is valid.")

      # 检查 ONNX 模型的输出类型
      for o in onnx_model.graph.output:
      print(o.name, o.type)

      # 使用 ONNX Runtime 进行推理
      ort_session = onnxruntime.InferenceSession(onnx_model_path)

      # 创建示例输入
      example_input = np.random.randn(1, 1, 240, 320).astype(np.float32)

      # 推理
      outputs = ort_session.run(
      None, # 输出名称(None 表示返回所有输出)
      {"input": example_input} # 输入名称和数据
      )

      # 打印输出形状
      print("semi shape:", outputs[0].shape) # semi 的形状
      print("desc shape:", outputs[1].shape) # desc 的形状
    3. 完成。

  • 等等。

(可选)在 C++项目中加载和推理 ONNX 模型

  • 在你的 C++ 项目中,引入 ONNX Runtime C++ API,并使用它来加载和推理 ONNX 模型。
  • 确保编译你的项目时,正确链接 ONNX Runtime 所需的库。

示例代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
#include <onnxruntime/core/providers/cpu/cpu_provider_factory.h>
#include <onnxruntime/core/session/onnxruntime_cxx_api.h>
#include <iostream>
#include <vector>
#include <random>

int main() {
// 1. 初始化 ONNX Runtime 环境
Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ONNXRuntime");

// 2. 创建 SessionOptions 并启用 CPU 提供器
Ort::SessionOptions session_options;
session_options.SetIntraOpNumThreads(1); // 设置线程数
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_BASIC);
Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CPU(session_options, 1));

// 3. 加载 ONNX 模型
const char* model_path = "./weights/superpoint.onnx";
Ort::Session session(env, model_path, session_options);
std::cout << "Successfully loaded ONNX model: " << model_path << std::endl;

// 4. 获取模型输入输出信息
size_t num_input_nodes = session.GetInputCount();
size_t num_output_nodes = session.GetOutputCount();

// 输入信息
Ort::AllocatorWithDefaultOptions allocator;
char* input_name = session.GetInputName(0, allocator);
Ort::TypeInfo input_type_info = session.GetInputTypeInfo(0);
auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
ONNXTensorElementDataType input_type = input_tensor_info.GetElementType();
std::vector<int64_t> input_dims = input_tensor_info.GetShape();

std::cout << "Input Name: " << input_name << std::endl;
std::cout << "Input Dimensions: ";
for (const auto& dim : input_dims) {
std::cout << dim << " ";
}
std::cout << std::endl;

// 输出信息
char* output_name_0 = session.GetOutputName(0, allocator);
char* output_name_1 = session.GetOutputName(1, allocator);
std::cout << "Output Names: " << output_name_0 << ", " << output_name_1 << std::endl;

// 5. 创建输入数据(随机 1x1x240x320 浮点数据)
std::vector<float> input_tensor_values(1 * 1 * 240 * 320);
std::generate(input_tensor_values.begin(), input_tensor_values.end(), []() -> float {
return static_cast<float>(rand()) / RAND_MAX;
});

std::vector<int64_t> input_shape = {1, 1, 240, 320}; // NCHW 格式

// 6. 将输入数据封装为 ONNX Tensor
Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
memory_info,
input_tensor_values.data(),
input_tensor_values.size(),
input_shape.data(),
input_shape.size()
);

// 7. 推理并获取输出
std::vector<const char*> output_names = {output_name_0, output_name_1};
auto output_tensors = session.Run(
Ort::RunOptions{nullptr},
&input_name,
&input_tensor,
1,
output_names.data(),
output_names.size()
);

// 8. 解析输出
// Semi 输出
auto* semi_data = output_tensors[0].GetTensorMutableData<float>();
auto semi_shape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape();
std::cout << "Semi Shape: ";
for (const auto& dim : semi_shape) {
std::cout << dim << " ";
}
std::cout << std::endl;

// Desc 输出
auto* desc_data = output_tensors[1].GetTensorMutableData<float>();
auto desc_shape = output_tensors[1].GetTensorTypeAndShapeInfo().GetShape();
std::cout << "Desc Shape: ";
for (const auto& dim : desc_shape) {
std::cout << dim << " ";
}
std::cout << std::endl;

// 释放资源
allocator.Free(input_name);
allocator.Free(output_name_0);
allocator.Free(output_name_1);

return 0;
}

优化 ONNX 模型

模型量化

注意事项

  1. 操作支持问题:确保 PyTorch 模型的所有算子(operator)都被 ONNX 支持。如果遇到不兼容的算子,可以尝试升级 ONNX 或 PyTorch,或者手动重写模型的部分代码以适配 ONNX。
  2. 性能优化:在导出 ONNX 模型时,可以启用 do_constant_folding 选项来执行常量折叠优化。此外,部署时可以使用 ONNX Runtime 的硬件加速(如 GPU)。
  3. 测试和验证:确保导出的 ONNX 模型在 C++ 项目中的推理结果与 PyTorch 模型一致。
  4. 在将一个 PyTorch 模型转换为 ONNX 模型时,不能直接在模型中使用 NumPy 的操作。如果某些功能确实需要 NumPy 而 PyTorch 不支持,可以尝试以下方法:
    • 将 NumPy 操作迁移到 PyTorch 实现中(用 PyTorch 提供的等价操作替代 NumPy 操作)。
    • 在 ONNX 导出后,使用自定义脚本对 ONNX 模型进行后处理,在部署阶段实现 NumPy 操作的功能。

模型量化

每次优化后,务必进行准确性和性能验证。

概述

模型量化(Model Quantization)是深度学习和机器学习中常用的一种模型压缩和加速技术。它的核心思想是将模型中原本用高精度(比如 32 位浮点数 float32)表示的权重和激活值,转换为低精度(比如 8 位整数 int8、16 位浮点数 float16、甚至更低)进行存储和计算。

主要内容和目标

  1. 压缩模型体积
    由于低精度数据类型占用更少的存储空间,量化后模型的文件体积会变小,有利于模型在移动端、嵌入式设备等存储受限的场景部署。
  2. 提高计算速度
    低精度数据类型运算速度更快,尤其是在支持低精度运算的硬件(如部分 CPU、GPU、NPU 等)上,可以显著提升推理速度。
  3. 降低能耗
    低精度运算通常功耗更低,非常适合边缘设备和移动端。

量化方式

  1. 动态量化(Dynamic Quantization)
    在模型推理阶段动态地将部分权重/激活转换为低精度。
  2. 静态量化(Static Quantization)
    在模型部署前,利用校准数据分析分布,提前将权重/激活转换为低精度。
  3. 量化感知训练(QAT, Quantization Aware Training)
    在模型训练阶段就模拟量化效果,使模型在低精度下表现更优。.

举例说明

假设有一个模型参数是:0.123456789 (float32, 32 位)

量化后转为 int8 表示:12 (int8, 8 位)

推理时再用缩放因子把 int8 恢复到原始范围。

ONNX 官方工具链优化

ONNX Simplifier

工具:onnx-simplifier

作用:简化模型结构,去除冗余节点,合并算子。

使用示例:

  • 安装:

    1
    2
    3
    conda activate your_env
    pip3 install -U pip
    pip3 install onnxsim
  • 命令行:

    1
    2
    onnxsim input_onnx_model output_onnx_model
    onnxsim -h # 如需更多高级功能,请尝试以下命令获取帮助信息
  • Python API:

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    import onnx
    from onnxsim import simplify

    # load your predefined ONNX model
    model = onnx.load(filename)

    # convert model
    model_simp, check = simplify(model)

    assert check, "Simplified ONNX model could not be validated"

    # use model_simp as a standard ONNX model object

    您可以在 onnxsim/onnx_simplifier.py 中查看 API 的更多详细信息

  • 等等。

ONNX Optimizer

工具:onnx/optimizer

作用:通过一系列内置 passes 对模型图进行结构优化(如常量折叠、冗余算子去除等)。

使用示例:

  • 安装:

    1
    2
    3
    conda activate your_env
    pip3 install -U pip
    pip3 install onnxoptimizer
  • 命令行:

    1
    2
    python3 -m onnxoptimizer -h  # 参数列表
    python -m onnxoptimizer input_model.onnx output_model.onnx
  • Python API:

    1
    2
    from onnx import optimizer
    optimized_model = optimizer.optimize(model, passes=['fuse_bn_into_conv'])
  • 等等。

运行时推理引擎优化

ONNX Runtime Graph Optimization

  • ONNX Runtime 默认自带丰富的图优化能力(如算子融合、常量折叠等)。

  • 你可以设置不同的优化等级(ORT_DISABLE_ALL, ORT_ENABLE_BASIC, ORT_ENABLE_EXTENDED, ORT_ENABLE_ALL)。

  • 示例:

    1
    2
    3
    4
    import onnxruntime as ort
    sess_options = ort.SessionOptions()
    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    ort.InferenceSession("model.onnx", sess_options)
  • 等等。

ONNX Runtime Quantization

工具:onnxruntime.quantization

作用:支持静态/动态/量化感知训练等多种量化方式,降低模型计算与存储需求。

示例(动态量化):

1
2
from onnxruntime.quantization import quantize_dynamic, QuantType
quantize_dynamic('model.onnx', 'model_quant.onnx', weight_type=QuantType.QInt8)

ONNX Runtime TensorRT / OpenVINO / CUDA 等后端加速

  • 利用硬件加速插件(如 TensorRT、OpenVINO 等)进一步提升推理速度。
  • 如果目标平台是 NVIDIA GPU,可以使用 TensorRT 对 ONNX 模型进行优化。TensorRT 可以自动进行层融合、精度校正、内存优化等操作,以减小模型大小并提高推理速度。

图变换与算子融合

  • 手动优化模型结构,如合并卷积与 BN、移除无用节点、精简分支结构等。
  • 可以通过 ONNX GraphSurgeon(NVIDIA 出品,适合 TensorRT 前处理)等工具编辑模型图。

其它常用优化方法

  • 剪枝(Pruning):通过去掉部分连接或节点实现模型瘦身(需配合训练/微调)。
  • 混合精度(FP16)/低精度(INT8):减少存储和计算量,提升推理效率。
  • Batch Normalization 融合:将 BN 层与前面的 Conv 层合并,减少推理步骤。
  • 常量折叠(Constant Folding):将可以提前计算的表达式直接固化为常量。

ONNX Runtime C++ API

概述

ONNX Runtime 是一个高性能的推理(Inference)引擎,专门用于运行基于 ONNX(Open Neural Network Exchange)格式的机器学习模型。它由微软(Microsoft)开发和维护,旨在提供跨平台、跨硬件的高效模型推理解决方案。

主要特点

  • 跨平台:支持 Windows、Linux、macOS、Android、iOS 等多种操作系统。
  • 高性能:针对 CPU、GPU(NVIDIA CUDA、DirectML 等)、FPGA、ARM 设备等多种硬件进行了优化,加速模型推理。
  • 多语言支持:提供 Python、C/C++、C#、Java、JavaScript 等多种语言的 API。
  • 兼容性强支持多种主流深度学习框架(如 PyTorch、TensorFlow)导出的 ONNX 模型。
  • 灵活部署:可用于云端、边缘端和本地环境,适合不同场景下的 AI 应用部署。

典型用途

  • 模型推理:在生产环境或终端设备上高效运行训练好的 ONNX 模型,实现实时预测和推断。
  • 跨框架迁移可以将不同深度学习框架训练的模型转为 ONNX 格式,通过 ONNX Runtime 部署到不同平台。
  • 硬件加速:充分利用不同平台的硬件能力,实现更低延迟和更高吞吐量的 AI 服务。

tips

  • GetTensorData() 用于读取张量数据,保证数据不可修改,适合只读操作。
  • GetTensorMutableData() 用于修改张量数据,适合需要直接操作张量内容的场景。

TensorRT

TensorRT 是 NVIDIA 开发的高性能深度学习推理库,专为 NVIDIA GPU 优化。

主要作用

  • 模型优化与加速:将深度学习模型(尤其是 ONNX/PyTorch/TensorFlow)进行图优化、层融合、精度削减(FP32→FP16/INT8)后,生成适合 NVIDIA GPU 推理的高效序列(engine)。
  • 极致推理性能:大幅提升模型在 NVIDIA 显卡上的推理速度(低延迟、高吞吐),适用于自动驾驶、边缘计算、云 AI 等场景。
  • 支持量化与混合精度:支持 FP32、FP16、INT8 等多种精度,兼顾速度和精度需求。
  • 与 DL 框架协作良好:可直接加载 ONNX 等格式的模型,简化集成过程。

常用流程

  1. 导入模型:支持ONNX、Caffe、TensorFlow等格式。
  2. 优化与编译:自动进行图结构优化(如层融合、张量合并等)、量化等。
  3. 生成引擎:编译为engine文件(特定硬件可用的二进制)。
  4. 推理部署:调用TensorRT API、C++/Python接口进行推理。

ONNX与TensorRT的异同

相同点

  • 都用于“模型部署”环节,提升训练模型从实验到实际产品生产环境的效率和性能。
  • 都支持模型推理加速,特别是配合NVIDIA硬件。
  • 常在同一个工程pipeline下配合使用。例如训练好PyTorch模型→ONNX导出→TensorRT优化→GPU部署。

不同点

特性 ONNX TensorRT
类型 通用开源的模型交换格式/标准 特定平台的推理加速引擎
作用 训练到部署的“格式桥梁”;跨框架、平台兼容 NVIDIA GPU推理优化和执行
支持硬件 通用CPU、GPU、FPGA、TPU、NPU、各种厂商芯片 仅NVIDIA GPU (含Jetson TX, Xavier等)
输入模型格式 PyTorch、TensorFlow等主流都可转导出 支持ONNX、Caffe、TensorFlow等
优化能力 取决于下游runtime,如ONNX Runtime等 定制为NVIDIA硬件极致优化
推理引擎/运行时 ONNX自有Runtime、OpenVINO等 TensorRT专属引擎
精度支持 取决于runtime,有些支持FP16/INT8 默认支持FP32、FP16、INT8混合精度
社区生态 标准格式、广泛支持、第三方工具生态极其丰富 主要NVIDIA生态,工业部署极多
  • ONNX已成为深度学习模型训练-部署-迁移的工业标准格式,在模型迁移到异构硬件或适配不同推理框架时广泛应用。
  • TensorRT则是面向NVIDIA硬件推理加速的主流、权威工具,是工业界优化推理性能的重要手段,尤其在需要低延迟/高吞吐/大批量推理的场景。
  • 两者不是“互斥”关系,很多生产系统正是把ONNX和TensorRT配合起来,既得到跨平台灵活性,又得到顶级推理性能。

关系示意

  1. PyTorch/TensorFlow训练 → 导出ONNX模型
  2. ONNX模型可用ONNX Runtime等通用推理引擎运行
  3. 也可用TensorRT载入ONNX模型,做极致NVIDIA GPU优化,生成engine文件并部署。
flowchart LR
    subgraph 训练端
        A[PyTorch/TensorFlow/Caffe]
        B[导出模型]
    end
    subgraph 部署端
        C[ONNX Runtime等推理引擎]
        D[TensorRT优化/编译]
        E[NVIDIA GPU引擎推理]
    end

    A --> B
    B -- ONNX --> C
    B -- ONNX --> D
    D --> E

ONNX知识点
http://zeyulong.com/posts/8b56b9a/
作者
龙泽雨
发布于
2025年8月7日
许可协议