PyTorch模型在无PyTorch环境下的部署：利用ONNX实现跨平台推理

1. 理解部署挑战与ONNX解决方案

在机器学习模型开发中，pytorch因其灵活性和易用性而广受欢迎。然而，当模型需要部署到资源受限或对依赖有严格要求的生产环境时，直接包含完整的pytorch库可能不切实际。这通常是由于库体积庞大、安装复杂性或与现有系统架构不兼容等原因。在这种场景下，我们需要一种机制，能够将训练好的pytorch模型“解耦”出来，使其能够在没有pytorch环境的情况下独立运行。

开放神经网络交换（ONNX, Open Neural Network Exchange）标准应运而生，它提供了一种通用的、跨框架的模型表示格式。ONNX允许开发者将模型从一个深度学习框架（如PyTorch、TensorFlow）导出，然后在另一个框架或专门的推理引擎中加载和运行。ONNX的核心优势在于：

• 框架无关性： 摆脱特定框架的依赖。
• 高效推理： ONNX Runtime等推理引擎针对不同硬件平台进行了优化。
• 跨平台： 支持多种操作系统和编程语言。

因此，将PyTorch模型导出为ONNX格式，是解决在无PyTorch环境下部署模型问题的理想方案。

2. PyTorch模型导出到ONNX

PyTorch提供了内置的API来方便地将模型导出为ONNX格式。这个过程通常涉及以下几个关键步骤：

1. 加载或定义模型： 确保您有一个已训练好或结构完整的PyTorch模型实例。
2. 准备一个虚拟输入： ONNX导出过程需要一个示例输入张量来跟踪模型的计算图。这个虚拟输入的形状和数据类型必须与模型实际接收的输入一致。
3. 调用torch.onnx.export： 使用PyTorch提供的torch.onnx.export函数进行导出。

以下是一个详细的导出示例：

import torch
import torch.nn as nn

# 1. 定义一个简单的PyTorch模型作为示例
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(10, 5) # 输入特征10，输出特征5
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(5, 2)  # 输入特征5，输出特征2 (例如，二分类)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 实例化模型并加载预训练权重（如果需要）
model = SimpleNet()
# model.load_state_dict(torch.load('your_model_weights.pth')) # 如果有预训练权重
model.eval() # 设置为评估模式，禁用Dropout和BatchNorm等

# 2. 准备一个虚拟输入张量
# 假设模型期望的输入是 (batch_size, input_features)
# 这里我们使用 batch_size=1，input_features=10
dummy_input = torch.randn(1, 10)

# 3. 定义ONNX导出参数
onnx_file_path = "simple_net.onnx"
input_names = ["input"]
output_names = ["output"]

# 如果您的模型需要支持动态批处理大小，可以设置dynamic_axes
# 例如：{ 'input' : {0 : 'batch_size'}, 'output' : {0 : 'batch_size'} }
dynamic_axes = {
    'input' : {0 : 'batch_size'},    # 第0维（batch_size）是动态的
    'output' : {0 : 'batch_size'}
}

# 4. 执行ONNX导出
try:
    torch.onnx.export(
        model,                      # 待导出的模型
        dummy_input,                # 虚拟输入
        onnx_file_path,             # ONNX模型保存路径
        verbose=False,              # 是否打印导出详细信息
        input_names=input_names,    # 输入节点的名称
        output_names=output_names,  # 输出节点的名称
        dynamic_axes=dynamic_axes,  # 定义动态输入/输出维度
        opset_version=11            # ONNX操作集版本，建议使用较新的稳定版本
    )
    print(f"模型已成功导出到 {onnx_file_path}")
except Exception as e:
    print(f"模型导出失败: {e}")

关键参数说明：

• model: 要导出的PyTorch模型实例。
• args: 一个元组或张量，表示模型的示例输入。它用于追踪模型的计算图。
• f: 输出ONNX文件的路径。
• input_names和output_names: 用于为ONNX图中的输入和输出节点命名，这在后续推理时很有用。
• dynamic_axes: 这是一个字典，用于指定哪些输入/输出维度可以是动态的（例如，批处理大小）。这对于部署时需要灵活输入尺寸的模型至关重要。
• opset_version: 指定ONNX操作集版本。选择一个合适的版本很重要，因为它会影响支持的操作和兼容性。

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

一旦模型被导出为ONNX格式，就可以使用ONNX Runtime进行推理。ONNX Runtime是一个高性能的推理引擎，支持多种编程语言（Python, C++, C#, Java等）和硬件平台。

无阶未来模型擂台/AI 应用平台

无阶未来模型擂台/AI 应用平台，一站式模型+应用平台

35

查看详情

以下是使用Python和ONNX Runtime进行推理的示例：

import onnxruntime as ort
import numpy as np

# 1. 加载ONNX模型
onnx_file_path = "simple_net.onnx"
try:
    # 创建ONNX Runtime会话
    sess = ort.InferenceSession(onnx_file_path)
    print(f"ONNX模型 {onnx_file_path} 已成功加载。")
except Exception as e:
    print(f"ONNX模型加载失败: {e}")
    exit()

# 获取模型输入和输出的名称
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name
print(f"模型输入名称: {input_name}")
print(f"模型输出名称: {output_name}")

# 2. 准备推理输入数据
# 注意：输入数据需要是NumPy数组，并且数据类型要与模型期望的一致（通常是float32）
# 假设模型期望的输入是 (batch_size, 10)
# 这里我们使用 batch_size=2 来演示动态批处理
input_data = np.random.rand(2, 10).astype(np.float32)

# 3. 执行推理
try:
    # 构建输入字典
    inputs = {input_name: input_data}

    # 运行推理
    outputs = sess.run([output_name], inputs)

    # outputs是一个列表，包含所有输出张量
    result = outputs[0]

    print(f"推理结果形状: {result.shape}")
    print(f"部分推理结果:\n{result[:5]}") # 打印前5个结果
except Exception as e:
    print(f"ONNX模型推理失败: {e}")

ONNX Runtime推理步骤：

1. 安装ONNX Runtime： pip install onnxruntime
2. 创建InferenceSession： 加载ONNX模型文件。
3. 获取输入/输出名称： 通过sess.get_inputs()和sess.get_outputs()获取模型输入和输出节点的名称。
4. 准备输入数据： 将您的数据转换为NumPy数组，并确保数据类型与模型期望的匹配（通常是np.float32）。
5. 运行推理： 调用sess.run()方法，传入输出名称列表和输入字典。

对于C++等其他语言的部署，ONNX Runtime也提供了相应的API。例如，在C++项目中，您可以包含ONNX Runtime的头文件，链接其库，然后使用Ort::Env、Ort::Session等类进行模型加载和推理。如果您的Python应用程序需要与C++进行交互（如原问题中提到的PyBind11），可以在C++部分使用ONNX Runtime，并通过PyBind11封装C++的推理函数，供Python调用。

4. 注意事项与最佳实践

在将PyTorch模型导出到ONNX并进行部署时，需要注意以下几点：

• 模型兼容性： 并非所有PyTorch操作都能直接映射到ONNX。复杂的自定义层或不常见的操作可能需要在PyTorch中进行修改或使用自定义ONNX算子。导出后，务必使用ONNX工具（如Netron）检查导出的ONNX图结构。
• 模型验证： 导出ONNX模型后，强烈建议在同一组输入数据上比较PyTorch模型和ONNX模型（通过ONNX Runtime）的输出，确保数值上的一致性。微小的差异可能是由于浮点精度或操作实现差异造成的。
• 动态输入： 如果模型需要处理可变大小的输入（例如，不同批次大小或不同图像分辨率），请务必在torch.onnx.export时正确配置dynamic_axes。
• model.eval()： 在导出之前，始终将PyTorch模型设置为评估模式（model.eval()）。这会禁用诸如Dropout和BatchNorm等在训练和推理阶段行为不同的层，确保推理结果的确定性。
• opset_version： 选择一个稳定的ONNX opset_version。过旧的版本可能不支持新的操作，过新的版本可能在某些推理引擎中尚未完全支持。
• 性能优化： ONNX Runtime支持多种执行提供者（Execution Providers），如CUDA、TensorRT、OpenVINO等，可以针对特定硬件进行优化，显著提升推理性能。在部署时，根据目标硬件环境选择合适的执行提供者。
• 模型大小： ONNX模型通常比包含整个PyTorch框架的部署方案更小，但仍需注意模型本身的参数量，以满足部署环境的存储和内存限制。

5. 总结

通过将PyTorch模型导出为ONNX格式，我们能够有效地解决在无PyTorch依赖环境中部署模型的挑战。ONNX提供了一个标准的、跨框架的模型表示，结合ONNX Runtime等高效推理引擎，使得PyTorch模型能够以轻量级、高性能的方式集成到各种生产系统中。遵循上述导出和推理的最佳实践，可以确保模型的顺利部署和稳定运行。

以上就是PyTorch模型在无PyTorch环境下的部署：利用ONNX实现跨平台推理的详细内容，更多请关注php中文网其它相关文章！