4、模型转换与优化工具链：ONNX、TensorRT、OpenVINO、TFLite、RKNN

模型训练出来只是第一步，怎么把它塞进嵌入式设备里跑起来，才是真正的硬仗。我见过太多团队，模型在PC上跑得飞起，一到板子上就卡成PPT。说白了，中间缺的就是这一环——模型转换与优化工具链。

今天咱们就聊聊市面上主流的几个工具链：ONNX、TensorRT、OpenVINO、TFLite、RKNN。每个我都踩过坑，有些坑现在想起来还肉疼。

4.1 ONNX：模型转换的“通用语言”

ONNX（Open Neural Network Exchange）是微软和Meta牵头搞的。它的定位很明确——做模型格式的“翻译官”。

你想想看，PyTorch训练出来的模型是.pt，TensorFlow是.pb，Keras是.h5。这些格式互相不认，怎么办？ONNX就是中间那个桥梁。

核心思路：先把模型导出成ONNX格式，再从这个通用格式转到目标平台的专用格式。

4.1.1 安装与配置

安装其实很简单，一行命令搞定：

pip install onnx onnxruntime

我个人习惯在虚拟环境里装，避免跟其他项目打架。如果你用的是PyTorch，还需要装个辅助包：

pip install onnx-simplifier

这个工具很有用。我遇到过好几次，导出的ONNX模型里有多余的节点，用simplifier一清理，模型体积直接小了一半。

4.1.2 模型导出流程

以PyTorch为例，导出ONNX的代码长这样：

import torch
import torch.onnx

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

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

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    'my_model.onnx',
    opset_version=11,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'},
                  'output': {0: 'batch_size'}}
)

这里有个坑——opset_version。我刚开始做的时候，随手写了个13，结果目标平台的推理引擎不支持。后来我学乖了，先查目标平台支持哪个版本，再决定。

小技巧：导出后一定要用 onnx.checker.check_model() 验证一下。我曾经导出一个模型，看起来没问题，一跑就崩，查了半天发现是某个算子版本不兼容。

4.2 TensorRT：NVIDIA的加速利器

如果你用的是NVIDIA的GPU（比如Jetson系列），那TensorRT基本是绕不开的。它能把模型优化到极致——FP16、INT8量化，层融合，内存复用，一套组合拳下来，推理速度能翻好几倍。

我记得第一次在Jetson Nano上跑一个分类模型，没优化前是30ms，用了TensorRT之后直接降到8ms。当时我盯着屏幕愣了好几秒。

4.2.1 安装与配置

TensorRT的安装稍微麻烦点。它不像ONNX那样pip一下就行。你需要去NVIDIA官网下载对应的deb包或者tar包。

以Jetson平台为例：

# 查看JetPack版本
cat /etc/nv_tegra_release

# 安装TensorRT（通常JetPack已经预装了）
sudo apt-get install tensorrt

# 验证安装
dpkg -l | grep TensorRT

如果是x86的Linux，建议用tar包解压，然后配置环境变量：

export LD_LIBRARY_PATH=/path/to/TensorRT/lib:$LD_LIBRARY_PATH
export PATH=/path/to/TensorRT/bin:$PATH

注意：TensorRT的版本跟CUDA版本是绑定的。我吃过这个亏——装了TensorRT 8.5，结果CUDA是11.4，死活跑不起来。后来查文档才发现，TensorRT 8.5要求CUDA 11.6以上。

4.2.2 模型转换流程

从ONNX转到TensorRT，有两种方式：

  1. 使用trtexec命令行工具——适合快速验证
  2. 使用Python API——适合集成到自动化流程

命令行方式最简单：

trtexec --onnx=my_model.onnx \
        --saveEngine=my_model.trt \
        --fp16 \
        --workspace=1024

Python API方式更灵活：

import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)
parser = trt.OnnxParser(network, TRT_LOGGER)

with open('my_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
config.set_flag(trt.BuilderFlag.FP16)

serialized_engine = builder.build_serialized_network(network, config)
with open('my_model.trt', 'wb') as f:
    f.write(serialized_engine)

避坑指南：我曾经在转换一个分割模型时，发现输出结果全是黑的。排查了半天，原来是TensorRT默认的输入输出格式跟ONNX不一致。解决办法是在转换时显式指定输入输出的数据类型。

4.3 OpenVINO：Intel平台的王牌

如果你用的是Intel的CPU、集成显卡或者Movidius神经计算棒，那OpenVINO就是你的首选。它能把模型优化到充分利用Intel硬件的特性。

说实话，我第一次用OpenVINO是在一个工业检测项目上。客户指定要用Intel的CPU，不能加GPU。我当时心里直打鼓——CPU跑深度学习，能行吗？结果OpenVINO一上，推理速度从200ms降到了40ms。嗯，真香。

4.3.1 安装与配置

OpenVINO的安装方式也很多。我推荐用pip安装，最省事：

pip install openvino-dev

如果你需要模型优化工具，还得装：

pip install openvino-dev[onnx,pytorch,tensorflow]

安装完后，记得初始化环境：

source /opt/intel/openvino_2023/setupvars.sh

4.3.2 模型转换流程

OpenVINO用mo（Model Optimizer）工具做转换。从ONNX转过来：

mo --input_model my_model.onnx \
   --output_dir ./openvino_model \
   --data_type FP16 \
   --input_shape [1,3,224,224]

转换后会生成两个文件：.xml（模型结构）和.bin（权重）。

加载推理的代码：

from openvino.runtime import Core

core = Core()
model = core.read_model('openvino_model/my_model.xml')
compiled_model = core.compile_model(model, 'CPU')

# 推理
input_data = ...  # 你的输入数据
result = compiled_model([input_data])

个人经验：OpenVINO对某些自定义算子支持不太好。我遇到过模型里有torch.where，OpenVINO直接报错。解决办法是改模型结构，或者用ONNX的custom_op机制。

4.4 TFLite：移动端和边缘设备的标配

TFLite是TensorFlow Lite的简称，专门为移动设备和嵌入式设备设计的。它支持Android、iOS、Linux、MCU等各种平台。

我最早接触TFLite是在一个手机APP项目上。模型要跑在手机上，还不能太耗电。TFLite的量化功能帮了大忙——模型体积从50MB压缩到12MB，速度还快了3倍。

4.4.1 安装与配置

pip install tensorflow

TFLite是TensorFlow的一部分，装好TensorFlow就能用。如果你只需要TFLite运行时，可以装轻量版：

pip install tflite-runtime

4.4.2 模型转换流程

从TensorFlow模型转TFLite：

import tensorflow as tf

# 加载模型
model = tf.saved_model.load('my_saved_model')

# 转换
converter = tf.lite.TFLiteConverter.from_saved_model('my_saved_model')
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]

tflite_model = converter.convert()
with open('my_model.tflite', 'wb') as f:
    f.write(tflite_model)

从ONNX转TFLite稍微麻烦点，需要先转成TensorFlow，再转TFLite。或者用onnx2tf工具：

pip install onnx2tf
onnx2tf -i my_model.onnx -o tflite_model

注意：TFLite的量化对精度有影响。我做过一个检测模型，INT8量化后mAP掉了3个点。后来用了混合量化（部分层用FP16，部分层用INT8），才把精度拉回来。

4.5 RKNN：Rockchip平台的专属优化

RKNN是Rockchip（瑞芯微）推出的神经网络推理框架。如果你用的是RK3588、RK3568这些芯片，那RKNN就是你的不二之选。

我在一个边缘计算盒子上用过RK3588，配合RKNN，跑YOLOv5s能做到30fps以上。说实话，这个性价比确实可以。

4.5.1 安装与配置

RKNN工具链分两部分：PC端的模型转换工具，和板端的运行时库。

PC端安装：

# 下载RKNN Toolkit
git clone https://github.com/rockchip-linux/rknn-toolkit2.git
cd rknn-toolkit2
pip install -r requirements.txt
pip install package/rknn_toolkit2-1.5.0-cp38-cp38-linux_x86_64.whl

板端需要交叉编译或者用预编译的库。我建议直接用Rockchip提供的Debian系统，里面已经预装了。

4.5.2 模型转换流程

从ONNX转RKNN：

from rknn.api import RKNN

rknn = RKNN()

# 配置
rknn.config(mean_values=[[123.675, 116.28, 103.53]],
            std_values=[[58.395, 57.12, 57.375]],
            target_platform='rk3588')

# 加载ONNX模型
ret = rknn.load_onnx(model='my_model.onnx')
if ret != 0:
    print('加载ONNX模型失败')
    exit(-1)

# 构建RKNN模型
ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
if ret != 0:
    print('构建RKNN模型失败')
    exit(-1)

# 导出
rknn.export_rknn('./my_model.rknn')

# 释放资源
rknn.release()

避坑指南：RKNN对模型中的某些算子支持有限。我曾经转一个带有torch.nn.functional.grid_sample的模型，RKNN直接报不支持。后来我手动把这个算子替换成了双线性插值，才搞定。

4.6 工具链对比与选型建议

工具链 目标平台 支持量化 安装难度 性能表现
ONNX 通用 有限 简单 中等
TensorRT NVIDIA GPU FP16/INT8 中等 优秀
OpenVINO Intel CPU/GPU FP16/INT8 简单 良好
TFLite 移动端/嵌入式 FP16/INT8 简单 良好
RKNN Rockchip芯片 INT8 中等 优秀

选型建议其实很简单：

  • 有NVIDIA GPU → 用TensorRT，别犹豫
  • Intel平台 → OpenVINO，白嫖Intel的优化
  • 手机或树莓派 → TFLite，生态最成熟
  • Rockchip芯片 → RKNN，原厂优化最到位
  • 做原型验证 → 先用ONNX，再转其他格式

好了，这一章的内容就到这。下一章咱们聊聊量化——怎么把模型从FP32压到INT8，精度还不掉太多。这可是个技术活。

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