PyTorch模型导出ONNX：在无PyTorch环境中高效推理

import torch
import torch.nn as nn
import numpy as np

# 定义一个简单的模型
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)

# 实例化模型并加载预训练权重（此处简化为随机初始化）
model = SimpleModel()
# 实际应用中，这里会加载训练好的模型权重，例如：
# model.load_state_dict(torch.load('path/to/your/model_weights.pth'))
model.eval() # 切换到评估模式，这对于导出ONNX至关重要，因为它会禁用Dropout等训练特有的层

# 准备一个虚拟输入张量，用于追踪模型计算图
# 这个虚拟输入的形状和数据类型必须与模型的实际输入匹配
dummy_input = torch.randn(1, 10) # 批大小为1，输入特征为10的张量

# 定义ONNX模型的保存路径
onnx_path = "MLmodel.onnx"

# 导出模型到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'], # 定义输入张量的名称
                       output_names=['output_tensor'],# 定义输出张量的名称
                       dynamic_axes={'input_tensor': {0: 'batch_size'},    # 声明输入张量的批次维度是动态的
                                     'output_tensor': {0: 'batch_size'}})   # 声明输出张量的批次维度是动态的
    print(f"模型已成功导出到 {onnx_path}")
except Exception as e:
    print(f"模型导出失败: {e}")

import onnxruntime as ort
import numpy as np

# ONNX模型的路径
onnx_path = "MLmodel.onnx"

try:
    # 创建ONNX Runtime会话
    # providers参数可以指定运行时使用的执行提供者，例如'CPUExecutionProvider'或'CUDAExecutionProvider'
    # 默认情况下，ONNX Runtime会尝试使用可用的最优化提供者。
    session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])

    # 获取模型的输入和输出名称
    # ONNX Runtime的输入和输出信息存储在session.get_inputs()和session.get_outputs()中
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    print(f"模型输入名称: {input_name}")
    print(f"模型输出名称: {output_name}")

    # 准备输入数据
    # 输入数据必须是NumPy数组，并且数据类型（如np.float32）和形状要与ONNX模型期望的匹配
    # 假设模型的输入是 (batch_size, 10)
    A = np.random.rand(1, 10).astype(np.float32) # 单个样本，10个特征，数据类型为float32

    print(f"输入数据形状: {A.shape}, 类型: {A.dtype}")

    # 执行推理
    # session.run()方法接收一个输出名称列表和一个输入字典
    results = session.run([output_name], {input_name: A})
    Result = results[0] # ONNX Runtime返回一个列表，通常我们取第一个元素作为结果

    print("推理结果:", Result)

except Exception as e:
    print(f"ONNX Runtime推理失败: {e}")