深度学习模型在C++平台的部署
一、概述深度学习模型能够在各种生产场景中发挥重要的作用,而深度学习模型往往在Python环境下完成训练,因而训练好的模型如何在生产环境下实现稳定可靠的部署,便是一个重要内容。C++开发平台广泛存在于各种复杂的生产环境,随着业务效能需求的不断提高,充分运用深度学习技术的优势显得尤为重要。本文介绍如何实现将深度学习模型部署在C++平台上。
二、步骤
s1. Python环境中安装深度学习框架(如PyTorch、TensorFlow等);
s2. P ython环境中设计并训练深度学习模型;
s3. 将训练好的模型保存为.onnx格式的模型文件;
s4. C++环境中安装Microsoft.ML.OnnxRuntime程序包;
(Visual Studio 2022中可通过项目->管理NuGet程序包完成快捷安装)
s5. C++环境中加载模型文件,完成功能开发。
三、示例
在Python环境下设计并训练一个关于手写数字识别的卷积神经网络(CNN)模型,将模型导出为ONNX格式的文件,然后在C++环境下完成对模型的部署和推理。
1. Python训练和导出
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.functional import F
# 定义简单的CNN模型
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
self.fc1 = nn.Linear(32 * 7 * 7, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 32 * 7 * 7)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
# 数据预处理
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
# 加载训练数据
train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
# 初始化模型、损失函数和优化器
model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 训练模型
def train(model, train_loader, criterion, optimizer, epochs=5):
model.train()
for epoch in range(epochs):
running_loss = 0.0
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}')
# 训练模型
train(model, train_loader, criterion, optimizer)
# 导出为ONNX格式
dummy_input = torch.randn(1, 1, 28, 28)
torch.onnx.export(
model,
dummy_input,
"mnist_model.onnx",
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)
print("模型已成功导出为mnist_model.onnx")
2. C++ 部署和推理
#include #include #include #include int main() { // 初始化环境 Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "MNIST"); Ort::SessionOptions session_options; session_options.SetIntraOpNumThreads(1); session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); // 加载模型 std::wstring model_path = L"mnist_model.onnx"; Ort::Session session(env, model_path.c_str(), session_options); // 准备输入 std::vector input_shape = { 1, 1, 28, 28 }; size_t input_tensor_size = 28 * 28; std::vector input_tensor_values(input_tensor_size); // 读取测试图片 cv::Mat test_image = cv::imread("test.jpg", cv::IMREAD_GRAYSCALE); // 将Mat数据复制到vector中 for (int i = 0; i < test_image.rows; ++i) { for (int j = 0; j < test_image.cols; ++j) { input_tensor_values = static_cast(test_image.at(i, j)); // 注意:uchar是unsigned char的缩写,表示无符号字符,通常用于存储灰度值 } } // 创建输入张量 auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); Ort::Value input_tensor = Ort::Value::CreateTensor( memory_info, input_tensor_values.data(), input_tensor_size, input_shape.data(), 4); // 设置输入输出名称 std::vector input_names; std::vector output_names; input_names.push_back(session.GetInputNameAllocated(0, Ort::AllocatorWithDefaultOptions()).get()); output_names.push_back(session.GetOutputNameAllocated(0, Ort::AllocatorWithDefaultOptions()).get()); // 运行推理 auto output_tensors = session.Run( Ort::RunOptions{ nullptr }, input_names.data(), &input_tensor, 1, output_names.data(), 1); // 获取输出结果 float* output = output_tensors.GetTensorMutableData(); std::vector results(output, output + 10); // 找到预测的数字 int predicted_digit = 0; float max_probability = results; for (int i = 1; i < 10; i++) { if (results > max_probability) { max_probability = results; predicted_digit = i; } } std::cout
页:
[1]