4、模型格式与转换：ONNX作为中间格式、PyTorch/TF模型转ONNX、ONNX转TensorRT/OpenVINO、量化与精度损失

模型格式转换，说白了就是给模型「换衣服」。你训练好的 PyTorch 或 TensorFlow 模型，边缘设备根本不认识。你得把它变成设备能懂的格式。

我个人习惯把 ONNX 当作中间桥梁。为什么？因为它生态好，几乎所有推理引擎都支持。我在项目中遇到过好几次，团队用不同框架训练模型，最后全靠 ONNX 统一收口。

ONNX 作为中间格式

ONNX 全称 Open Neural Network Exchange，微软和 Meta 联合搞的。它的核心价值就一句话：一次转换，到处推理。

你想想看，如果每个框架都直接对接 TensorRT、OpenVINO、TFLite，那得写多少适配代码？ONNX 帮你把 N 对 N 的问题，简化成了 N 对 1 对 N。

ONNX 的核心优势：

  • 算子标准统一，跨框架兼容
  • 支持静态图和动态图导出
  • 内置量化、优化工具链
  • 社区活跃，更新及时

嗯，这里要注意：ONNX 不是万能的。有些 PyTorch 的自定义算子，ONNX 不支持。我踩过这个坑，后来老老实实改模型结构。

PyTorch 模型转 ONNX

PyTorch 转 ONNX 用 torch.onnx.export。代码不长，但参数有讲究。

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    export_params=True,
    opset_version=11,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里我重点说几个参数：

  • opset_version：算子集版本。我建议用 11 或更高，太低有些新算子不支持。
  • dynamic_axes：动态轴。如果你要支持可变 batch size，必须配这个。
  • do_constant_folding：常量折叠。能优化掉一些固定计算，建议打开。

我的经验：导出前一定要调用 model.eval()。我见过有人忘了这步，结果导出的 ONNX 里还带着 dropout 和 BN 的训练状态，推理结果全乱套。

TensorFlow 模型转 ONNX

TF 转 ONNX 用 tf2onnx 工具。支持 SavedModel、Keras H5、Checkpoint 等多种格式。

import tf2onnx
import tensorflow as tf

model = tf.keras.applications.MobileNetV2(weights='imagenet')
model.save('mobilenetv2_savedmodel')

# 命令行方式
# python -m tf2onnx.convert --saved-model mobilenetv2_savedmodel --output mobilenetv2.onnx

# Python API 方式
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"))
output_path = "mobilenetv2.onnx"

model_proto, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    opset=13,
    output_path=output_path
)

说实话，TF 转 ONNX 比 PyTorch 麻烦一点。主要是 TF 的算子太「灵活」，有些动态控制流 ONNX 根本接不住。我曾经遇到一个 TF 模型里用了 tf.while_loop，转 ONNX 直接报错。最后只能重写模型。

避坑指南：TF 2.x 的 Keras 模型导出时，注意 input_shape 要明确指定。我遇到过 batch size 没写，结果导出的 ONNX 推理时 batch 维度被固定成 1。

ONNX 转 TensorRT

TensorRT 是 NVIDIA 的推理优化引擎。ONNX 转 TensorRT，核心是 trtexec 工具或 Python API。

# 命令行方式
trtexec --onnx=resnet18.onnx \
        --saveEngine=resnet18.engine \
        --fp16 \
        --workspace=1024 \
        --minShapes=input:1x3x224x224 \
        --optShapes=input:8x3x224x224 \
        --maxShapes=input:16x3x224x224

# Python API 方式
import tensorrt as trt

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

with open("resnet18.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)

profile = builder.create_optimization_profile()
profile.set_shape("input", (1, 3, 224, 224), (8, 3, 224, 224), (16, 3, 224, 224))
config.add_optimization_profile(profile)

engine = builder.build_serialized_network(network, config)
with open("resnet18.engine", "wb") as f:
    f.write(engine)

这里有个关键点：动态 shape 的配置。minShapes、optShapes、maxShapes 三个值必须配好。我见过有人只配了固定 shape，结果换不同分辨率图片时直接崩了。

性能对比（ResNet-18，batch=1，FP32 vs FP16）：

引擎 延迟（ms） 吞吐量（fps） 显存占用（MB）
ONNX Runtime（CPU） 12.3 81 —
TensorRT FP32 3.1 322 245
TensorRT FP16 1.8 555 128

你看，FP16 比 FP32 快了将近一倍，显存还省一半。这就是 TensorRT 的价值。

ONNX 转 OpenVINO

OpenVINO 是 Intel 的推理引擎，主要跑在 CPU、集成显卡和 VPU 上。转换用 mo 工具（Model Optimizer）。

# 命令行方式
mo --input_model resnet18.onnx \
   --output_dir ./openvino_model \
   --data_type FP16 \
   --input_shape [1,3,224,224] \
   --mean_values [123.675,116.28,103.53] \
   --scale_values [58.395,57.12,57.375]

# 转换后生成两个文件：
# resnet18.xml  - 模型结构
# resnet18.bin  - 模型权重

OpenVINO 有个好处：对 CPU 优化极好。我曾在 Intel Xeon 上跑过，比 ONNX Runtime 快 2-3 倍。但要注意，它不支持所有 ONNX 算子。我遇到过 ScatterND 算子不支持，最后用 --disable_fusing 参数绕过去的。

我的建议：如果目标设备是 Intel CPU 或集成显卡，优先用 OpenVINO。如果是 NVIDIA GPU，用 TensorRT。别混着用，效果反而差。

量化与精度损失

量化，说白了就是把 FP32 的模型变成 INT8。模型变小了，推理变快了，但精度会掉。

我刚开始做量化时，觉得精度损失无所谓。结果有一次在安防项目里，量化后的模型把行人误检成车辆，差点出事。从那以后，我对量化精度特别敏感。

常见的量化方式有三种：

  • 训练后量化（PTQ）：最简单，拿少量校准数据跑一遍，自动算量化参数。精度损失一般 1-3%。
  • 量化感知训练（QAT）：训练时模拟量化效果，精度损失可以控制在 0.5% 以内。但需要重新训练，成本高。
  • 动态量化：只量化权重，激活值运行时动态量化。适合 LSTM、Transformer 这类模型。

量化精度对比（ResNet-18，ImageNet Top-1）：

精度类型 Top-1 准确率 模型大小 推理速度（fps）
FP32 69.76% 44.7 MB 322
INT8（PTQ） 68.12% 11.2 MB 890
INT8（QAT） 69.31% 11.2 MB 890

你看，PTQ 掉了 1.64%，QAT 只掉了 0.45%。如果你的业务对精度要求高，别省那点训练成本，上 QAT。

避坑指南：量化校准数据一定要覆盖真实场景。我曾经用 ImageNet 的校准集量化了一个安防模型，结果部署后精度暴跌。后来发现，安防场景的图片分布和 ImageNet 完全不一样。换了真实场景的校准数据后，精度才恢复正常。

最后说一句：模型格式转换和量化，不是一锤子买卖。你很可能需要反复调参、测试、再调参。我建议你建一个自动化 pipeline，把转换、量化、精度验证串起来。这样每次改模型，跑一遍 pipeline 就知道效果。

嗯，这一节就到这里。下一节我们聊聊推理引擎的部署实战，到时候我会拿一个完整的项目案例来演示。

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