模型导出与格式转换：ONNX、TensorRT、OpenVINO等格式介绍，PyTorch模型转ONNX实战

模型训练完了，只是万里长征走了一半。另一半，就是把它部署到目标硬件上。这一步，说白了就是让模型「说人话」——让不同框架、不同硬件都能听懂。

我见过太多团队，训练时跑得飞起，一到部署就抓瞎。为什么？因为训练框架和推理引擎之间，缺了一座桥。今天我们就来聊聊这座桥怎么搭。

为什么需要中间格式？

你想想看，PyTorch训练出来的模型，本质上是Python对象。它依赖Python运行时、依赖CUDA、依赖各种第三方库。但你的目标设备可能是个手机、是个边缘盒子、甚至是个MCU。它跑不了Python，更装不了PyTorch。

所以我们需要一个「中间人」——一个通用的、轻量的、硬件友好的模型表示。这就是ONNX、TensorRT、OpenVINO这些格式存在的意义。

三大主流格式速览

格式	维护方	核心优势	适用场景
ONNX	微软、Meta等	框架互通、生态广	跨平台、模型交换
TensorRT	NVIDIA	GPU极致优化	NVIDIA显卡推理
OpenVINO	Intel	CPU/VPU优化	Intel硬件部署

我个人习惯是：先用ONNX做「通用桥接」，再根据目标硬件做二次优化。比如跑在NVIDIA卡上，就转TensorRT；跑在Intel CPU上，就转OpenVINO。

ONNX：模型界的「普通话」

ONNX全称是Open Neural Network Exchange。说白了，它定义了一套标准化的计算图描述方式。不管你是PyTorch、TensorFlow还是PaddlePaddle，都能导出成ONNX。反过来，任何支持ONNX的推理引擎都能加载它。

我在项目中遇到过一个问题：团队里有人用PyTorch，有人用TensorFlow。模型互相转换时，精度对不上。后来统一用ONNX作为中间格式，问题就解决了。嗯，这就是标准化的力量。

TensorRT：NVIDIA的「性能怪兽」

如果你用的是NVIDIA显卡，TensorRT几乎是绕不开的选择。它做了几件事：层融合、精度校准（FP16/INT8）、内存复用、内核自动调优。效果嘛——我见过一个模型，用TensorRT跑INT8，比原生PyTorch快了5倍，精度只掉了0.3%。

但要注意，TensorRT对某些算子支持不完善。比如动态shape、某些自定义算子，可能会报错。我曾经花了一整天，就为了把一个PyTorch的Gather操作改成TensorRT支持的版本。

OpenVINO：Intel生态的「万金油」

OpenVINO是Intel推出的推理框架。它主要优化CPU和集成显卡（核显）上的推理性能。如果你部署在Intel的服务器或边缘设备上，OpenVINO是个好选择。

它支持从ONNX、TensorFlow、PyTorch直接导入模型。我个人觉得它的Model Optimizer工具挺好用，能自动做图优化和量化。

PyTorch模型转ONNX实战

好了，理论说完了。我们来点实际的。下面这段代码，演示了如何把一个训练好的PyTorch模型导出为ONNX。

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("ONNX导出成功！")

这段代码看起来简单，但有几个坑要注意。

第一，模型一定要设成eval模式。否则BatchNorm和Dropout的行为会不一样，导出的ONNX可能精度不对。我刚开始做的时候忘了这步，结果推理结果和训练时对不上，排查了半天。

第二，dummy_input的shape要和实际输入一致。如果你模型输入是224x224，你给个128x128，导出不会报错，但推理时就会出问题。

第三，dynamic_axes这个参数很关键。它告诉ONNX：这个维度是可变的。比如batch_size，你训练时是1，但部署时可能想一次跑8张图。如果不设dynamic_axes，导出的ONNX就固定了batch_size=1，没法改。

import onnxruntime as ort
import numpy as np

# 加载ONNX模型
session = ort.InferenceSession("resnet18.onnx")

# 准备输入
input_name = session.get_inputs()[0].name
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 推理
outputs = session.run(None, {input_name: input_data})

# 对比PyTorch结果
with torch.no_grad():
    torch_output = model(torch.from_numpy(input_data))

print("ONNX结果:", outputs[0][0][:5])
print("PyTorch结果:", torch_output[0][:5])

如果两者结果差异在1e-5以内，说明导出没问题。如果差异很大，就要检查模型结构或算子兼容性了。

常见问题与避坑

• 算子不支持：某些PyTorch操作（如torch.einsum）在ONNX中没有对应算子。解决办法是改写模型，用ONNX支持的算子替代。
• 动态shape问题：如果你的模型输入尺寸不固定，导出时一定要设dynamic_axes。否则推理时输入尺寸变了就会报错。
• 精度损失：FP16或INT8量化后，精度可能会下降。建议先做精度验证，再决定是否使用低精度。

我曾经遇到过一个案例：一个语义分割模型，导出ONNX后推理结果全是黑的。排查了半天，发现是模型里用了一个自定义的插值操作，ONNX不支持，直接给忽略了。后来改用了torch.nn.functional.interpolate，问题解决。

总结

模型导出这件事，说白了就是「翻译」。把PyTorch的「方言」翻译成ONNX的「普通话」，再根据目标硬件翻译成TensorRT或OpenVINO的「本地话」。每一步都可能遇到坑，但只要掌握了方法，就能稳稳地走通。

下一章，我们会深入TensorRT的优化技巧。到时候我会分享一些实战中的调优经验，包括如何做INT8量化、如何避免精度损失。嗯，敬请期待。