4、模型转换与量化：TensorRT、ONNX Runtime、TFLite 的部署流程，FP16/INT8 量化实战

好，咱们进入正题。模型训练完了，精度看着也不错，但离真正跑在嵌入式设备上，还差最关键的一步——部署。说白了，就是把训练好的“大胖子”模型，变成能在目标硬件上高效运行的“精干小伙”。这一步，我称之为“模型落地前的最后一公里”。

今天咱们就聊聊三个主流的部署工具：TensorRT、ONNX Runtime 和 TFLite。它们各有各的脾气，也各有各的绝活。我会结合 FP16 和 INT8 量化，把整个流程拆开揉碎了讲给你听。

4.1 为什么需要模型转换与量化？

你想想看，训练时我们用 PyTorch 或 TensorFlow，模型里存的参数都是 FP32 的。FP32 精度高，但体积大、计算慢。在嵌入式设备上，内存和算力都金贵得很。

模型转换，就是把模型从训练框架的格式，转成目标推理引擎能识别的格式。比如从 .pt 转成 .onnx，或者 .engine。

量化，则是把 FP32 的权重和激活值，用更低位宽的数据类型来表示。比如 FP16（16位浮点）或者 INT8（8位整型）。

核心收益：

  • 体积缩小：FP32 转 INT8，模型体积直接缩到原来的 1/4。
  • 速度提升：INT8 计算比 FP32 快 2-4 倍，在专用硬件上甚至更快。
  • 功耗降低：更少的计算量，意味着更低的功耗，这对电池供电的设备至关重要。

嗯，这里要注意，量化不是免费的午餐。精度会有损失，关键看你怎么平衡。

4.2 TensorRT：NVIDIA 的看家本领

如果你用的是 NVIDIA 的 Jetson 系列，或者带 GPU 的嵌入式板卡，那 TensorRT 几乎是绕不开的选择。我个人习惯，只要项目里用了 NVIDIA 的硬件，推理引擎首选就是 TensorRT。

TensorRT 的牛逼之处在于，它不只是做量化，还会做图优化、层融合、内核自动调优。说白了，它把你的模型从头到脚重新编排了一遍，让 GPU 跑得最舒服。

4.2.1 转换流程

从 PyTorch 模型到 TensorRT 引擎，我一般走这条路：

  1. 导出 ONNX：先把 PyTorch 模型转成 ONNX 格式。这是中间桥梁。
  2. 构建 TensorRT 引擎：用 TensorRT 的 Python API 或者 trtexec 工具，读取 ONNX 模型，指定精度（FP32、FP16 或 INT8），然后构建引擎。
  3. 序列化与反序列化：构建好的引擎可以保存成 .engine 文件。部署时直接加载，省去每次重新构建的时间。

这里有个坑，我曾经踩过。ONNX 导出时，如果模型里有动态尺寸的输入，比如 batch size 不固定，需要显式指定动态轴。否则 TensorRT 构建时会报错。

避坑指南：我曾经在导出 ONNX 时忘了设置 dynamic_axes，结果 TensorRT 构建死活过不去。后来查了半天才发现，ONNX 默认是静态尺寸。记得在 torch.onnx.export 里加上 dynamic_axes 参数。

4.2.2 FP16 量化实战

FP16 量化是最简单的。在构建 TensorRT 引擎时，指定精度为 FP16 即可。几乎不需要额外操作，精度损失通常很小。

import tensorrt as trt

# 创建 builder 和 network
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 解析 ONNX 模型
parser = trt.OnnxParser(network, TRT_LOGGER)
with open(onnx_model_path, 'rb') as f:
    parser.parse(f.read())

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

# 关键：启用 FP16
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)

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

你看，核心就一行代码：config.set_flag(trt.BuilderFlag.FP16)。但要注意，不是所有 GPU 都支持 FP16 加速。Jetson 系列基本都支持，但老一点的桌面 GPU 可能不行。

4.2.3 INT8 量化实战

INT8 量化就复杂多了。它需要一个“校准数据集”，用来统计激活值的分布，从而找到最佳的量化参数。

我一般准备几百张有代表性的图片，不需要标注，只要覆盖模型可能遇到的各种场景就行。

import tensorrt as trt
from calibrator import EntropyCalibrator2  # 自定义校准器

# 创建校准器
calibrator = EntropyCalibrator2(calibration_data, batch_size=32)

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

# 启用 INT8
if builder.platform_has_fast_int8:
    config.set_flag(trt.BuilderFlag.INT8)
    config.int8_calibrator = calibrator

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

我的经验：校准数据集的质量直接影响 INT8 量化后的精度。我建议用 500-1000 张图片，覆盖不同光照、角度、背景。如果校准集太单一，量化后的模型在复杂场景下可能会崩。

4.3 ONNX Runtime：跨平台的瑞士军刀

如果你的目标平台不是 NVIDIA，或者你想在 CPU 上跑，那 ONNX Runtime 是个好选择。它支持多种硬件后端，包括 CPU、GPU、甚至一些 NPU。

ONNX Runtime 的部署流程相对简单。你只需要把模型转成 ONNX 格式，然后直接用 ONNX Runtime 加载推理就行。

4.3.1 转换与推理

import onnxruntime as ort

# 创建推理会话
session = ort.InferenceSession("model.onnx")

# 获取输入输出名称
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 准备输入数据
import numpy as np
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 推理
result = session.run([output_name], {input_name: input_data})

你看，代码非常简洁。但 ONNX Runtime 的量化，需要借助 onnxruntime.quantization 工具包。

4.3.2 量化实战

ONNX Runtime 支持静态量化和动态量化。静态量化需要校准数据，动态量化不需要，但精度稍差。

from onnxruntime.quantization import quantize_static, QuantType
from onnxruntime.quantization import CalibrationMethod

# 静态量化
quantized_model = quantize_static(
    model_input="model.onnx",
    model_output="model_quantized.onnx",
    calibration_data_reader=calibration_data_reader,
    quant_format=QuantType.QInt8,
    per_channel=True,
    calibration_method=CalibrationMethod.MinMax
)

注意：ONNX Runtime 的 INT8 量化，在 CPU 上加速效果明显，但在 GPU 上可能不如 TensorRT。我建议在 CPU 部署时优先考虑 ONNX Runtime 量化。

4.4 TFLite：移动端和边缘设备的王者

如果你的目标是手机、树莓派、或者微控制器，那 TFLite 是首选。它专门为资源受限的设备做了优化。

TFLite 的转换流程，我一般用 TensorFlow 的 TFLiteConverter。

4.4.1 转换流程

import tensorflow as tf

# 加载 TensorFlow 模型
model = tf.saved_model.load("saved_model")

# 创建转换器
converter = tf.lite.TFLiteConverter.from_saved_model("saved_model")

# 设置优化选项
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# 设置代表数据集（用于量化）
def representative_dataset():
    for _ in range(100):
        data = np.random.randn(1, 224, 224, 3).astype(np.float32)
        yield [data]

converter.representative_dataset = representative_dataset

# 指定目标精度
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

# 转换
tflite_model = converter.convert()

# 保存
with open("model_quantized.tflite", "wb") as f:
    f.write(tflite_model)

4.4.2 FP16 与 INT8 选择

TFLite 支持 FP16 和 INT8 量化。FP16 量化后模型体积减半，精度损失极小。INT8 量化体积减到 1/4，但精度损失稍大。

量化类型 体积缩减 速度提升 精度损失 适用场景
FP16 ~50% ~1.5x 极小 对精度要求高的场景
INT8 ~75% ~2-4x 中等 对速度要求高的场景

我的建议：如果设备支持硬件加速（比如手机上的 NPU），优先用 INT8。如果只是 CPU 推理，FP16 可能更稳妥。我一般在项目初期先用 FP16 跑通，再尝试 INT8 看精度能否接受。

4.5 量化后的精度验证

量化不是终点，验证才是。我每次量化完，都会跑一遍完整的测试集，对比量化前后的精度差异。

如果精度下降超过 1-2%，我会考虑：

  • 使用更大量的校准数据集
  • 尝试不同的校准方法（比如 KL 散度 vs 均方误差）
  • 对敏感层进行混合精度量化（部分层用 FP16，部分层用 INT8）

嗯，这里要提醒一句：不要盲目追求 INT8。如果 FP16 已经能满足性能要求，就别折腾 INT8 了。省下来的时间，多喝杯咖啡不香吗？

4.6 总结

模型转换与量化，是嵌入式视觉部署的核心技能。TensorRT 适合 NVIDIA 平台，ONNX Runtime 适合跨平台 CPU 部署，TFLite 适合移动端和边缘设备。

FP16 量化简单高效，INT8 量化需要更多调优。我个人习惯是：先用 FP16 快速验证，再根据精度和速度需求决定是否上 INT8。

记住，部署不是一锤子买卖。你可能需要反复调整量化参数、校准数据集，甚至修改模型结构，才能达到最佳效果。但当你看到模型在设备上流畅运行时，那种成就感，值得。

← 上一章 📖 返回目录 下一章 →