4. ONNX与模型转换：PyTorch/TensorFlow模型转ONNX、ONNX Runtime推理优化

模型训练出来只是第一步，真正要让它跑在田间的边缘设备上，中间还有一道绕不过去的坎——模型转换。说白了，就是把PyTorch或TensorFlow的模型，变成一种通用的、能被各种推理引擎识别的格式。ONNX就是目前最主流的中间格式。

我个人习惯把ONNX比作「模型界的普通话」。不管你的模型是用什么方言（框架）写的，只要转成ONNX，就能在支持ONNX的设备上跑起来。这章我们就来聊聊怎么把训练好的模型转成ONNX，以及怎么用ONNX Runtime做推理优化。

4.1 为什么需要ONNX？

你想想看，农业场景下的设备五花八门。有的用Jetson Nano，有的用树莓派，还有的用手机芯片。每个设备支持的推理框架都不一样。如果每次部署都要重新写一遍模型，那工程师得累死。

ONNX（Open Neural Network Exchange）解决了这个问题。它定义了一套统一的算子集和计算图格式。只要你的模型能转成ONNX，就能在支持ONNX Runtime、TensorRT、OpenVINO等引擎的设备上运行。

核心优势：

  • 一次转换，多处部署：ONNX格式可以在不同硬件和推理引擎间复用
  • 算子标准化：避免了框架间算子不兼容的问题
  • 推理优化：ONNX Runtime等引擎会对计算图做自动优化
  • 量化支持：ONNX原生支持INT8、FP16等低精度推理

4.2 PyTorch模型转ONNX

PyTorch转ONNX相对简单，因为PyTorch官方提供了torch.onnx.export接口。我在项目中遇到过最坑的事情是——动态轴没设置对，导致推理时输入尺寸变了就报错。

下面是一个标准的转换流程：

import torch
import torchvision.models as models

# 1. 加载模型
model = models.resnet18(pretrained=True)
model.eval()

# 2. 创建示例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出ONNX
torch.onnx.export(
    model,                      # 模型
    dummy_input,                # 示例输入
    "resnet18.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'}
    }
)

print("模型转换完成：resnet18.onnx")

避坑指南：

我曾经因为没设置dynamic_axes，导致部署时只能跑固定batch size的输入。后来在田间测试时，图片尺寸稍微一变就崩了。记住：只要你的推理场景可能用到不同尺寸的输入，一定要设置动态轴。

4.3 TensorFlow模型转ONNX

TensorFlow转ONNX稍微麻烦一点，因为TF的图结构和PyTorch的动态图不太一样。我推荐用tf2onnx这个工具，它支持TF 1.x和2.x。

先安装工具：

pip install tf2onnx onnx onnxruntime

然后执行转换：

import tensorflow as tf
import tf2onnx

# 1. 加载TF模型
model = tf.keras.applications.MobileNetV2(weights='imagenet')

# 2. 指定输入签名
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

# 3. 转换
output_path = "mobilenetv2.onnx"
model_proto, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    opset=13,
    output_path=output_path
)

print(f"转换完成：{output_path}")

注意：

TF转ONNX时，如果模型里用了自定义层或TF特有的算子（比如tf.image.resize），可能会转换失败。我建议先检查算子兼容性，必要时用ONNX支持的算子替换。

4.4 ONNX模型验证与可视化

转换完成后，别急着部署。先验证一下ONNX模型是否正确。我习惯用onnx.checker做完整性检查，再用netron可视化计算图。

import onnx
import onnxruntime as ort

# 1. 加载ONNX模型
onnx_model = onnx.load("resnet18.onnx")

# 2. 检查模型结构
onnx.checker.check_model(onnx_model)
print("模型结构检查通过")

# 3. 用ONNX Runtime做推理验证
ort_session = ort.InferenceSession("resnet18.onnx")

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

# 执行推理
outputs = ort_session.run(
    None,
    {"input": input_data}
)

print(f"推理输出shape：{outputs[0].shape}")

可视化的话，我推荐用Netron（https://netron.app），直接把ONNX文件拖进去就能看到完整的计算图。这在排查模型结构问题时特别有用。

4.5 ONNX Runtime推理优化

ONNX Runtime（ORT）是微软开源的推理引擎，专门为ONNX模型做了大量优化。说白了，它能把你的模型跑得更快、更省内存。

ORT的优化主要体现在几个方面：

优化类型 说明 效果
图优化 合并相邻算子、消除冗余计算 减少推理延迟10%-30%
算子融合 将多个小算子合并为一个大算子 减少内存访问次数
内存优化 复用中间张量、减少内存分配 降低峰值内存占用
量化 FP32转INT8/FP16 推理速度提升2-4倍
执行提供者 支持CUDA、TensorRT、OpenVINO等后端 充分利用硬件加速

下面是一个使用ORT进行推理优化的示例：

import onnxruntime as ort
import numpy as np
import time

# 1. 创建推理会话（开启所有优化）
options = ort.SessionOptions()
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
options.enable_profiling = True

# 2. 选择执行提供者（优先使用CUDA）
providers = [
    'CUDAExecutionProvider',
    'CPUExecutionProvider'
]

ort_session = ort.InferenceSession(
    "resnet18.onnx",
    sess_options=options,
    providers=providers
)

# 3. 准备输入
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 4. 预热（第一次推理通常较慢）
_ = ort_session.run(None, {"input": input_data})

# 5. 正式推理并计时
start = time.time()
for _ in range(100):
    outputs = ort_session.run(None, {"input": input_data})
end = time.time()

avg_time = (end - start) / 100 * 1000
print(f"平均推理时间：{avg_time:.2f} ms")

我的经验：

在农业场景下，我建议优先使用ORT_ENABLE_ALL级别的图优化。对于边缘设备，如果内存紧张，可以关闭enable_profiling来减少额外开销。另外，如果设备支持CUDA，一定要把CUDAExecutionProvider放在前面，GPU推理速度通常是CPU的5-10倍。

4.6 量化：让模型在边缘设备上跑得更快

量化是部署到边缘设备的关键技术。说白了，就是把模型参数从FP32压缩到INT8或FP16。精度会损失一点点，但推理速度能翻好几倍。

ONNX Runtime支持两种量化方式：

  • 动态量化：推理时动态将权重转为INT8，简单但效果一般
  • 静态量化：用校准数据集预先统计量化参数，效果更好

静态量化的示例：

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

# 准备校准数据（从训练集中取100张图片）
calibration_data = np.random.randn(100, 3, 224, 224).astype(np.float32)

# 执行静态量化
quantize_static(
    model_input="resnet18.onnx",
    model_output="resnet18_int8.onnx",
    calibration_data_reader=calibration_data,
    quant_format=QuantType.QInt8,
    per_channel=True,
    calibration_method=CalibrationMethod.MinMax
)

print("量化完成：resnet18_int8.onnx")

注意：

量化后的模型精度可能会下降1%-3%。我建议在量化前后用验证集做一次精度对比。如果精度下降太多，可以尝试只量化部分层，或者使用混合精度（部分层FP16，部分层INT8）。

4.7 本章知识体系

下面这张图总结了ONNX模型转换与优化的完整流程：

ONNX模型转换与优化流程 ONNX模型转换与优化流程 PyTorch模型 TensorFlow模型 其他框架模型 转换工具：torch.onnx.export / tf2onnx / onnx.convert ONNX模型（.onnx） 验证：onnx.checker / onnxruntime推理验证 / Netron可视化 优化：图优化 / 算子融合 / 量化（INT8/FP16）/ 执行提供者选择 核心思路：一次转换 → 多处部署 → 持续优化

嗯，到这里，ONNX模型转换与优化的核心内容就讲完了。从PyTorch/TensorFlow转ONNX，到验证、量化、推理优化，每一步都有坑，但也都有对应的解决方案。我个人觉得，掌握好ONNX这套工具链，是农业AI模型部署的必修课。毕竟，模型再好，跑不起来也是白搭。

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