第4章：模型基础与格式：ONNX、TensorRT与量化原理

各位同学，今天我们来聊聊模型格式和量化。说实话，这块内容在嵌入式部署里是绕不开的坎。我刚开始做部署的时候，也踩过不少坑，比如模型转了一半报错，或者量化后精度掉得没法看。今天我把这些经验都揉碎了讲给你听。

4.1 ONNX模型格式详解

ONNX，全称是Open Neural Network Exchange。说白了，它就是AI模型界的“通用语言”。你想想看，PyTorch训练出来的模型，TensorFlow训练出来的模型，它们各自有各自的格式。但嵌入式设备上跑的推理引擎，比如TensorRT、OpenVINO，它们不可能支持所有框架的原生格式。这时候，ONNX就派上用场了。

ONNX的核心价值：它定义了一套标准化的算子集和计算图表示方法。你的模型只要导出成ONNX格式，理论上就能在任何支持ONNX的推理引擎上运行。

我在项目中遇到过一个问题：用PyTorch训练好的分类模型，导出ONNX后在TensorRT上推理结果全错。后来排查发现，是PyTorch中某个自定义的激活函数，ONNX标准算子集里没有对应的映射。所以，导出ONNX时一定要检查算子兼容性。

4.1.1 ONNX的导出流程

以PyTorch为例，导出ONNX其实就几行代码：

import torch
import torch.onnx

# 假设你有一个训练好的模型
model = torch.load('model.pth')
model.eval()

# 创建一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 输入张量
    'model.onnx',        # 输出文件名
    export_params=True,  # 导出权重参数
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

嗯，这里要注意opset_version。版本越高，支持的算子越多，但兼容性可能变差。我个人习惯用opset 11或13，这两个版本比较稳定。

4.1.2 ONNX的检查与优化

导出后，我建议你用onnxruntime跑一遍推理，验证结果是否正确：

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession('model.onnx')
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
result = sess.run([output_name], {input_name: input_data})
print(result)

如果结果和原始模型一致，恭喜你，ONNX导出成功了。如果不一致，那就要检查算子映射或者输入输出的数据类型了。

4.2 TensorRT模型格式

TensorRT是NVIDIA推出的推理优化引擎。它专门针对NVIDIA GPU做了深度优化。说白了，它能把你的模型“编译”成一种高度优化的格式，推理速度能快好几倍。

TensorRT支持两种主要的模型格式：

• ONNX格式：直接加载ONNX模型，然后由TensorRT进行优化和编译。
• TensorRT引擎文件（.engine/.trt）：这是TensorRT编译后的二进制文件，包含了优化后的计算图和权重。加载后可以直接推理，不需要再编译。

我个人更推荐使用引擎文件。为什么呢？因为每次编译都需要时间，尤其是在嵌入式设备上，编译一次可能要几分钟。如果你把引擎文件保存下来，下次直接加载，就能省掉这个时间。

4.2.1 从ONNX到TensorRT引擎

转换流程大致如下：

import tensorrt as trt

# 创建TensorRT日志记录器
logger = trt.Logger(trt.Logger.WARNING)

# 创建构建器
builder = trt.Builder(logger)

# 创建网络定义
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)

# 创建ONNX解析器
parser = trt.OnnxParser(network, logger)

# 解析ONNX模型
with open('model.onnx', 'rb') as f:
    parser.parse(f.read())

# 创建构建配置
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

# 构建引擎
engine = builder.build_serialized_network(network, config)

# 保存引擎文件
with open('model.engine', 'wb') as f:
    f.write(engine)

这段代码看起来简单，但实际项目中我遇到过不少坑。比如set_memory_pool_limit设置得太小，会导致构建失败。我曾经在Jetson Nano上折腾了半天，最后发现是显存不够，把workspace调大就好了。

4.2.2 TensorRT的优化策略

TensorRT的优化主要体现在以下几个方面：

• 算子融合：把多个连续的算子合并成一个，减少计算开销。
• 精度校准：支持FP16和INT8量化，减少模型体积和计算量。
• 内存优化：复用内存缓冲区，减少内存分配和释放的开销。
• 内核自动调优：为每个算子选择最优的CUDA内核实现。

你想想看，一个卷积层后面跟着一个ReLU激活函数，在原始框架里是两个操作。但TensorRT会把它们融合成一个Conv+ReLU的CUDA内核，省掉了中间结果的读写。这就是优化的精髓。

4.3 模型量化原理（INT8/FP16）

量化，说白了就是把模型的权重和激活值从高精度（比如FP32）转换成低精度（比如FP16或INT8）。这样做的好处很明显：模型体积变小，推理速度变快，功耗降低。但代价是精度可能会下降。

我在项目中做过一个对比：一个FP32的ResNet-50模型，体积是98MB。量化成INT8后，体积变成了25MB，推理速度提升了3倍，但Top-1准确率只下降了0.5%。这个trade-off，在嵌入式设备上是非常值得的。

4.3.1 FP16量化

FP16量化是最简单的。它就是把FP32的数值映射到FP16的范围内。FP16的表示范围是大约±65504，精度大约是3.3位有效数字。对于大多数模型来说，FP16量化几乎不会造成精度损失。

在TensorRT中启用FP16很简单：

config.set_flag(trt.BuilderFlag.FP16)

就这么一行代码。但要注意，不是所有算子都支持FP16。如果某个算子不支持，TensorRT会自动回退到FP32。

4.3.2 INT8量化

INT8量化就复杂多了。它需要把FP32的数值映射到-128到127的整数范围内。这里的关键是缩放因子（scale）和零点（zero point）。

量化的数学公式很简单：

q = round(r / scale) + zero_point

其中，r是原始的FP32值，q是量化后的INT8值，scale是缩放因子，zero_point是零点偏移。

反量化就是反过来：

r = (q - zero_point) * scale

嗯，这里要注意，scale和zero_point的选取直接影响量化精度。选得不好，精度会掉得很厉害。

4.3.3 校准方法

为了找到合适的scale和zero_point，我们需要一个校准数据集。校准数据集不需要很大，几百张图片就够了。它的作用是统计模型各层激活值的分布，然后根据分布选择合适的量化参数。

常见的校准方法有三种：

方法	原理	特点
最大绝对值校准	取激活值的最大绝对值作为scale	简单，但对异常值敏感
KL散度校准	最小化量化前后分布的KL散度	精度较好，TensorRT默认使用
百分位校准	取某个百分位（如99.9%）的值作为scale	对异常值鲁棒

我个人习惯用KL散度校准。它在大多数场景下都能取得不错的精度。但如果你发现某个模型量化后精度掉得厉害，可以试试百分位校准。

4.3.4 量化感知训练（QAT）

如果后训练量化（PTQ）的精度不满足要求，可以考虑量化感知训练（QAT）。QAT在训练过程中模拟量化操作，让模型学会适应低精度表示。这样量化后的精度损失会更小。

QAT的流程大致如下：

1. 用FP32训练一个基线模型。
2. 在模型中插入伪量化节点（FakeQuantize）。
3. 用较小的学习率继续训练几个epoch。
4. 导出量化后的模型。

QAT的效果通常比PTQ好，但需要额外的训练时间和数据。在嵌入式项目中，如果精度要求很高，我建议用QAT。

4.4 本章小结

好了，今天的内容就到这里。我们讲了ONNX的导出和检查，TensorRT的引擎生成和优化策略，还有FP16和INT8量化的原理和实践。这些知识，说白了就是嵌入式部署的“基本功”。

你想想看，一个模型从训练到部署，中间要经过多少道工序？ONNX是桥梁，TensorRT是引擎，量化是加速器。这三样东西，缺一不可。

最后，送你一句话：模型部署不是简单的格式转换，而是一个系统工程。每一步都要细心，每一个参数都要理解。只有这样，你的模型才能在嵌入式设备上跑得又快又准。