PyTorch转ONNX部署实战教程

本文针对PyTorch模型部署难题，特别是当软件项目对依赖有严格限制，无法直接引入PyTorch运行时的情况，提供了一套实战指南。核心方案是将PyTorch模型转化为开放神经网络交换（ONNX）格式，实现模型与框架的解耦。通过`torch.onnx.export`函数，开发者可轻松将PyTorch模型导出为ONNX格式，并利用ONNX Runtime等推理引擎，在无PyTorch环境下高效执行模型推理。本文详细阐述了模型导出的关键参数设置，如`opset_version`、`dynamic_axes`等，并提供了Python和C++环境下的推理示例代码，助力开发者实现PyTorch模型的轻量化、跨平台部署，满足各种严苛的应用场景需求。

挑战：PyTorch模型在无PyTorch环境中的部署

在许多实际应用场景中，尤其是在资源受限、对依赖包有严格控制或需要跨平台部署的环境下，直接引入完整的PyTorch运行时是不可行或不理想的。例如，在嵌入式设备、移动应用、小型桌面软件或高性能推理服务中，开发者可能希望避免庞大的PyTorch依赖，转而寻求更轻量、更独立的模型部署方案。原始问题中提到的“小型的、最小依赖的软件”以及“开发者拒绝包含PyTorch”正是这一挑战的典型体现。

解决方案：利用ONNX实现模型独立部署

解决上述问题的关键在于将PyTorch训练好的模型转换为一种与框架无关、且能被多种推理引擎高效执行的标准格式。开放神经网络交换（ONNX, Open Neural Network Exchange）正是为此目的而生。ONNX是一个开放标准，定义了一套可扩展的计算图模型，支持不同框架之间的模型互操作性。通过将PyTorch模型导出为ONNX格式，我们可以在没有PyTorch的环境中，借助ONNX Runtime等专用推理引擎进行模型推理。

1. PyTorch模型导出为ONNX格式

将PyTorch模型导出为ONNX格式是一个相对直接的过程，主要通过torch.onnx.export函数完成。此过程需要一个“虚拟输入”（dummy input），用于追踪模型计算图的结构和张量形状。

示例代码：PyTorch模型导出

import torch
import torch.nn as nn

# 1. 定义一个简单的PyTorch模型（此处以一个简单的全连接网络为例）
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 2) # 假设输入特征维度为10，输出维度为2

    def forward(self, x):
        return self.fc(x)

# 2. 实例化模型并加载预训练权重（如果模型已训练）
model = SimpleModel()
# model.load_state_dict(torch.load('path/to/trained_weights.pth')) # 如果有预训练权重，请取消注释并加载
model.eval() # 将模型设置为评估模式，禁用Dropout和BatchNorm等层的训练行为

# 3. 创建一个虚拟输入（Dummy input）
# 虚拟输入的大小和类型必须与模型在实际推理时期望的输入一致。
# 例如，如果模型期望一个批次大小为1，特征维度为10的浮点张量：
dummy_input = torch.randn(1, 10) # batch_size=1, input_features=10

# 4. 定义ONNX模型保存路径
onnx_path = "simple_model.onnx"

# 5. 导出模型到ONNX格式
try:
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        export_params=True,        # 导出模型的所有参数
        opset_version=11,          # ONNX操作集版本，建议使用较新的稳定版本
        do_constant_folding=True,  # 是否执行常量折叠优化
        input_names=['input_tensor'],     # 定义ONNX图中输入节点的名称
        output_names=['output_tensor'],   # 定义ONNX图中输出节点的名称
        dynamic_axes={'input_tensor': {0: 'batch_size'},  # 允许输入批次大小动态变化
                      'output_tensor': {0: 'batch_size'}}
    )
    print(f"PyTorch模型已成功导出到 {onnx_path}")
except Exception as e:
    print(f"模型导出失败: {e}")

导出参数说明：

• model: 要导出的PyTorch模型实例。
• dummy_input: 一个用于追踪模型计算图的示例输入张量。其形状、数据类型和设备（CPU/GPU）应与实际推理时保持一致。
• f: ONNX模型文件的保存路径。
• export_params: 如果为True，则导出模型的所有参数（权重和偏置）。
• opset_version: 指定ONNX操作集版本。选择合适的版本以确保兼容性和功能支持。
• do_constant_folding: 启用常量折叠优化，有助于减小模型大小和提高推理效率。
• input_names和output_names: 为ONNX图中的输入和输出节点指定有意义的名称，方便后续推理时引用。
• dynamic_axes: 这是一个非常重要的参数，允许指定模型输入/输出张量的哪些维度可以是动态的（例如，批次大小或序列长度）。如果您的模型需要处理可变大小的输入，务必正确配置此参数。

2. 在无PyTorch环境下进行ONNX模型推理

一旦模型被导出为ONNX格式，就可以使用ONNX Runtime或其他兼容的推理引擎进行加载和推理，而无需PyTorch环境。ONNX Runtime是一个高性能的推理引擎，支持多种硬件平台和编程语言（如Python, C++, C#, Java等）。

示例代码：使用ONNX Runtime进行推理（Python）

import onnxruntime as ort
import numpy as np

# 1. 加载ONNX模型
onnx_path = "simple_model.onnx"
try:
    session = ort.InferenceSession(onnx_path)
    print(f"ONNX模型 {onnx_path} 已成功加载。")
except Exception as e:
    print(f"加载ONNX模型失败: {e}")
    exit()

# 2. 获取模型输入和输出名称
# 确保这些名称与导出时`input_names`和`output_names`中定义的名称一致
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
print(f"模型输入名称: {input_name}, 输出名称: {output_name}")

# 3. 准备输入数据（Numpy array）
# 注意：输入数据类型和形状必须与ONNX模型期望的一致。
# PyTorch通常使用torch.float32，对应Numpy的np.float32。
input_data = np.random.randn(1, 10).astype(np.float32) # 示例输入数据

# 4. 执行推理
# session.run的第一个参数是期望的输出名称列表，第二个参数是一个字典，
# 键是输入名称，值是对应的Numpy输入数据。
try:
    outputs = session.run([output_name], {input_name: input_data})

    # 5. 获取推理结果
    result = outputs[0]
    print("推理结果：", result)
    print("结果形状：", result.shape)
except Exception as e:
    print(f"ONNX模型推理失败: {e}")

对于C++环境：

对于需要在C++应用程序中进行推理的场景，可以使用ONNX Runtime C++ API。这通常涉及以下步骤：

1. 包含ONNX Runtime C++头文件。
2. 创建Ort::Env和Ort::Session对象来加载模型。
3. 准备输入张量（使用Ort::MemoryInfo和Ort::Value）。
4. 调用session.Run()执行推理。
5. 解析输出张量。

如果需要通过Python接口调用C++推理逻辑（如原问题中提到的PyBind11），则C++部分会负责加载ONNX模型并执行推理，PyBind11则负责将Python数据（如Numpy数组）桥接到C++，并以Python对象的形式返回推理结果。

注意事项

• model.eval()的重要性： 在导出模型之前，务必调用model.eval()。这会禁用Dropout层并冻结BatchNorm层的统计数据，确保模型在训练和推理时行为一致。
• 虚拟输入（Dummy Input）的准确性： 虚拟输入的形状、数据类型和设备（CPU/GPU）必须与模型实际推理时的输入完全匹配。不匹配可能导致导出失败或模型行为异常。
• opset_version的选择： 选择一个ONNX Runtime支持且能完整表达模型操作的opset_version。过旧的版本可能不支持某些操作，过新的版本可能尚未被广泛支持。
• 动态轴配置： 如果模型的输入或输出尺寸在推理时可能变化（例如，可变批次大小或图像尺寸），dynamic_axes参数的正确配置至关重要。
• 自定义操作符： 如果PyTorch模型中使用了自定义的PyTorch操作符，ONNX可能无法直接导出。在这种情况下，可能需要实现自定义ONNX操作符或寻找替代方案。
• 预处理和后处理： ONNX模型仅包含神经网络的计算图。图像的归一化、文本的Tokenization等预处理步骤，以及推理结果的解析、阈值判断等后处理步骤，通常需要在推理端独立实现。
• 模型验证： 导出后，建议使用ONNX Runtime加载导出的模型，并与原始PyTorch模型在相同输入下进行推理，比较输出结果，确保一致性。

总结

将PyTorch模型导出为ONNX格式是解决在无PyTorch依赖环境中部署模型的标准且高效的解决方案。这一方法不仅能够显著减少部署环境的依赖，实现模型的轻量化，还能利用ONNX Runtime等优化过的推理引擎，在多种硬件和操作系统上提供高性能的推理能力。通过遵循上述导出和推理步骤，并注意相关事项，开发者可以顺利地将PyTorch模型集成到对依赖有严格限制的软件项目中。

终于介绍完啦！小伙伴们，这篇关于《PyTorch转ONNX部署实战教程》的介绍应该让你收获多多了吧！欢迎大家收藏或分享给更多需要学习的朋友吧~golang学习网公众号也会发布文章相关知识，快来关注吧！