4、模型转换与优化：将PyTorch/TensorFlow模型导出为ONNX，使用onnx-simplifier和onnxoptimizer进行图优化

在多传感器融合推理项目中，PyTorch或TensorFlow训练出的模型往往包含大量仅服务于训练过程的算子、动态控制流以及冗余计算图结构。直接将这些模型部署到ONNX Runtime端侧，不仅会显著增加内存占用和启动延迟，还可能因算子兼容性问题导致推理失败。本章将系统讲解从训练框架导出ONNX的标准流程，并深入剖析如何利用onnx-simplifier与onnxoptimizer对计算图进行精简与优化，确保模型在端侧设备上以最小代价、最高效率运行。

4.1 模型导出的核心原理与前置准备

将PyTorch或TensorFlow模型导出为ONNX，本质上是将训练框架中的动态计算图“冻结”为静态的、与框架无关的中间表示。这一过程涉及三个关键层面：算子映射（将框架算子转换为ONNX标准算子）、权重固化（将训练得到的参数序列化为ONNX的初始器）、控制流扁平化（将Python层面的循环和条件分支转换为ONNX的Loop/If算子或直接展开）。对于多传感器融合模型，通常包含多个输入分支（图像、点云、雷达数据），导出时需要明确指定每个输入的名称、数据类型和动态维度。

4.2 PyTorch模型导出ONNX的标准流程

PyTorch通过torch.onnx.export接口实现模型导出。该接口基于追踪（Tracing）机制，即通过实际运行一次模型来记录算子执行序列。对于多传感器融合模型，需要构造一个包含所有输入分支的示例张量，并确保模型处于eval()模式且关闭梯度计算。

import torch
import torch.onnx

# 假设模型包含两个输入分支：图像(3, 480, 640)和点云(4, 128, 128)
model = MultiSensorFusionModel().eval()

# 构造示例输入
dummy_img = torch.randn(1, 3, 480, 640)
dummy_pc = torch.randn(1, 4, 128, 128)

# 导出ONNX
torch.onnx.export(
    model,
    (dummy_img, dummy_pc),
    "fusion_model.onnx",
    input_names=["image", "pointcloud"],
    output_names=["detection", "segmentation"],
    dynamic_axes={
        "image": {0: "batch"},
        "pointcloud": {0: "batch"},
        "detection": {0: "batch"},
        "segmentation": {0: "batch"}
    },
    opset_version=15,
    do_constant_folding=True
)
print("ONNX导出完成")

导出完成后，务必使用onnx.checker.check_model验证模型结构完整性，并使用onnx.shape_inference.infer_shapes进行形状推断，确保所有中间张量的维度信息明确。

4.3 TensorFlow模型导出ONNX的转换路径

TensorFlow模型导出ONNX通常通过tf2onnx工具完成。该工具支持从SavedModel、Keras H5或Checkpoint文件转换。对于包含预处理逻辑的模型，建议将预处理层（如归一化、Resize）一并包含在导出范围内，避免端侧额外实现。

import tf2onnx
import tensorflow as tf

# 加载Keras模型
model = tf.keras.models.load_model("fusion_model.h5")

# 指定输入签名
spec = (tf.TensorSpec((1, 480, 640, 3), tf.float32, name="image"),
        tf.TensorSpec((1, 128, 128, 4), tf.float32, name="pointcloud"))

# 转换并保存
onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=15)
with open("fusion_model.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())
print("TF转换ONNX完成")

4.4 onnx-simplifier：计算图简化与冗余消除

onnx-simplifier的核心目标是消除训练框架遗留的冗余计算，使计算图更简洁、更易于端侧优化。其主要功能包括：常量折叠（将仅依赖常量的子图预先计算）、形状推断与静态化（将动态形状转换为静态形状）、冗余节点删除（移除Identity、Dropout等训练专用节点）、算子融合（将连续的小算子合并为单个算子）。

使用onnx-simplifier有两种方式：命令行工具和Python API。推荐在Python脚本中集成，便于自动化处理。

import onnx
from onnxsim import simplify

# 加载原始模型
model = onnx.load("fusion_model.onnx")

# 执行简化
model_sim, check = simplify(model)
assert check, "简化后模型验证失败"

# 保存简化模型
onnx.save(model_sim, "fusion_model_sim.onnx")
print(f"简化完成，节点数从 {len(model.graph.node)} 降至 {len(model_sim.graph.node)}")

简化过程中，simplify函数会返回一个布尔值check，表示简化后的模型是否通过数值一致性验证。该验证通过随机输入比对原始模型与简化模型的输出差异，确保简化不改变模型语义。

4.5 onnxoptimizer：面向端侧的深度图优化

onnxoptimizer提供了更细粒度的图优化能力，其优化策略分为通用优化和特定后端优化两类。通用优化包括：算子融合（如Conv+BN融合、Gemm+Add融合）、内存布局优化（调整张量存储顺序以匹配端侧硬件）、冗余消除（删除无用的Reshape和Transpose）。特定后端优化则针对ONNX Runtime的CPU、GPU或NPU执行引擎进行定制。

import onnxoptimizer

# 加载简化后的模型
model = onnx.load("fusion_model_sim.onnx")

# 查看可用优化策略
all_passes = onnxoptimizer.get_available_passes()
print("可用优化策略:", all_passes)

# 执行全部优化
optimized_model = onnxoptimizer.optimize(model, passes=["fuse_consecutive_transposes", "eliminate_nop_dropout", "extract_constant_to_initializer"])

# 保存优化结果
onnx.save(optimized_model, "fusion_model_opt.onnx")
print("优化完成")

4.6 案例拆解：多传感器融合模型的完整优化闭环

某自动驾驶团队需要将包含Camera和LiDAR双分支的融合检测模型部署到Jetson Orin设备上。初始PyTorch模型包含120个节点，推理延迟为45ms，内存占用为680MB。团队按照以下流程进行优化：

在优化过程中，团队发现LiDAR分支中的自定义算子Voxelization无法被ONNX Runtime直接支持。解决方案是：将该算子替换为标准的ScatterND + Gather组合，并通过onnx-simplifier的custom_op接口注册自定义映射。经过调整后，模型在Jetson Orin上实现了12ms的推理延迟，满足实时性要求。

4.7 常见误区与验证策略

在模型转换与优化过程中，工程师常陷入以下误区：

• 过度依赖自动优化：onnx-simplifier和onnxoptimizer无法解决所有问题，特别是涉及自定义算子或复杂控制流时，需要手动修改模型结构。
• 忽略数值一致性验证：优化后的模型必须通过多组随机输入的输出比对，确保最大误差在可接受范围内（通常<1e-4）。
• 未考虑端侧算子支持：某些ONNX算子（如NonMaxSuppression）在端侧版本中可能不支持，需提前查阅ONNX Runtime的算子支持矩阵。

建议建立标准化的验证流程：第一步，使用onnxruntime的CPU版本进行数值比对；第二步，在目标端侧设备上运行基准测试，记录延迟和内存；第三步，使用真实传感器数据验证推理结果的语义正确性。只有通过全部验证的模型才能进入部署阶段。