win11下 “pytorch导出模型“ 以及 “C++使用onnxruntime部署”

部分一：PyTorch导出模型

在Win11下，PyTorch是一个强大的深度学习框架，它提供了丰富的工具来训练和导出模型。在这一部分，我们将使用鸢尾花数据集，演示如何在PyTorch中训练一个简单的模型，并将其导出为ONNX格式。

1、引言

深度学习模型的导出对于模型在不同平台上的部署至关重要。PyTorch的灵活性使得导出过程变得相对简单，同时保持了模型的准确性。

2、数据准备和模型训练

在这一步，我们首先加载鸢尾花数据集，对数据进行预处理，然后训练一个简单的神经网络模型。以下是代码示例：

# 导入所需的库
import torch
import torch.onnx
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# 加载鸢尾花数据集
iris = load_iris()
X, y = iris.data, iris.target

# 数据预处理
X_scaled = X / 10

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)

# 将数据转换为PyTorch的Tensor
X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train, dtype=torch.long)
X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
y_test_tensor = torch.tensor(y_test, dtype=torch.long)

# 定义简单的神经网络模型
class IrisModel(nn.Module):
    def __init__(self):
        super(IrisModel, self).__init__()
        self.fc1 = nn.Linear(4, 8)
        self.fc2 = nn.Linear(8, 3)  # 输入特征为4，输出类别为3

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.log_softmax(self.fc2(x), dim=1)
        return x

# 初始化模型、损失函数和优化器
model = IrisModel()
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(100):
    optimizer.zero_grad()
    output = model(X_train_tensor)
    loss = criterion(output, y_train_tensor)
    loss.backward()
    optimizer.step()

3、模型导出

在训练完成后，我们使用torch.onnx.export方法将模型导出为ONNX格式。导出的ONNX文件将在接下来的部署中使用，以下是代码示例：

# 将模型转换为ONNX格式
dummy_input = torch.randn(1, 4)  # 创建一个虚拟输入
onnx_path = 'iris_model.onnx'
torch.onnx.export(model, dummy_input, onnx_path,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes = {'input':
                            {0: 'batch_size'},
                       'output':
                            {0: 'batch_size'}
                       })

4、总结

在这一部分，我们演示了如何使用PyTorch训练一个简单的神经网络模型，并将其导出为ONNX格式，为模型在不同平台上的部署做好了准备。

部分二：C++使用ONNX Runtime部署

ONNX Runtime是一个用于高性能推理的开源引擎，它支持在不同平台上运行ONNX格式的模型。在这一部分，我们将学习如何使用C++和ONNX Runtime加载并运行先前导出的鸢尾花分类模型。

1、引言

ONNX Runtime的强大之处在于其跨平台性能，使得模型能够在各种设备上进行高效推理。

2、环境配置和项目设置

在使用C++部署模型之前，我们需要确保系统中已经正确安装了ONNX Runtime，并且我们的C++项目设置正确。
本文采用v1.16.3版本，下载地址：https://github.com/microsoft/onnxruntime/releases

3. C++代码实现

以下是一个简单的C++代码示例，演示如何加载ONNX模型并进行推理：

#include 
#include 
#include 
#include 

int main() {

    // ONNX模型文件路径
    const wchar_t* model_path = L"D:\\vs_project\\demo\\iris_model.onnx";

    // 创建ONNX运行环境和内存信息
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "Default");
    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);

    // 配置会话选项
    Ort::SessionOptions session_option;
    session_option.SetIntraOpNumThreads(5); // 设置并行线程数
    session_option.SetGraphOptimizationLevel(ORT_ENABLE_ALL);

    // 定义模型输入和输出的名称
    const char* input_names[] = { "input" };
    const char* output_names[] = { "output" };

    // 定义样本数量和输入输出矩阵的大小
    const int num_samples = 2;
    std::array<float, num_samples * 4> input_matrix;
    std::array<float, num_samples * 3> output_matrix;

    // 定义输入输出矩阵的形状
    std::array<int64_t, 2> input_shape{ num_samples, 4 };
    std::array<int64_t, 2> output_shape{ num_samples, 3 };

    // 定义样本输入数据
    std::vector<std::vector<float>> sample_x = { {2.1, 3.5, 1.4, 0.2}, {5.1, 1.5, 2.4, 0.2} };
    int sample_y = 3;

    // 将样本数据复制到输入矩阵中
    for (int i = 0; i < num_samples; i++) {
        for (int j = 0; j < 4; j++) {
            input_matrix[i * 4 + j] = sample_x[i][j];
        }
    }

    // 创建输入和输出Tensor
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_matrix.data(), input_matrix.size(), input_shape.data(), input_shape.size());
    Ort::Value output_tensor = Ort::Value::CreateTensor<float>(memory_info, output_matrix.data(), output_matrix.size(), output_shape.data(), output_shape.size());

    try {
        // 创建ONNX会话并运行模型
        Ort::Session session(env, model_path, session_option);
        session.Run(Ort::RunOptions{ nullptr }, input_names, &input_tensor, 1, output_names, &output_tensor, 1);

    }
    catch (const Ort::Exception& e) {
        // 处理ONNX Runtime异常
        std::cerr << "ONNX Runtime 异常: " << e.what() << std::endl;
    }
    catch (const std::exception& e) {
        // 处理标准异常
        std::cerr << "标准异常: " << e.what() << std::endl;
    }
    catch (...) {
        // 处理未知异常
        std::cerr << "未知异常." << std::endl;
    }

    // 输出预测结果
    std::cout << "--- 预测结果 ---" << std::endl;

    for (int i = 0; i < num_samples; i++) {
        std::cout << "输出矩阵: ";
        for (int j = 0; j < sample_y; j++) {
            std::cout << output_matrix[i * sample_y + j] << " ";
        }
        std::cout << std::endl;

        // 找到输出矩阵中的argmax值
        int argmax_value = std::distance(output_matrix.begin() + i * sample_y, std::max_element(output_matrix.begin() + i * sample_y, output_matrix.begin() + (i + 1) * sample_y));
        std::cout << "样本 " << i << " 的输出 argmax 值: " << argmax_value << std::endl;
    }

    // 等待用户按键结束程序
    getchar();
    return 0;
}

推理结果：

--- 预测结果 ---
输出矩阵: -0.0161782 -4.42303 -5.50902
样本 0 的输出 argmax 值: 0
输出矩阵: -2.30582 -0.799013 -0.797285
样本 1 的输出 argmax 值: 2

4. 总结

在这一部分，我们通过使用C++和ONNX Runtime，成功加载并运行了在PyTorch中训练并导出的鸢尾花分类模型。这为在不同C++支持的平台上进行模型推理提供了一个简单而强大的解决方案。

通过这两个部分，我们实现了从PyTorch训练模型到在C++环境中进行推理的全过程。这个流程可以在Win11下轻松实现，为模型的实际应用提供了一个完整的参考。