基于C++ ONNX的深度模型部署（以MNIST手写数字分类为例）Visual studio 2022开发

环境以及依赖

python下的环境除了必要的pytorch环境外，还需要安装onnx，直接输入命令pip install onnx。
Visual studio中需要安装onnxruntime环境，具体参考链接: link。版本的话我安装的是cpu版的（训练用gpu，推理可以使用cpu）。

代码构成

代码分为四部分，pytorch下的模型构建，pytorch下的模型训练，pytorch下的模型导出，C++环境下onnx模型的推理部署。

模型代码（python）

为了避免额外的麻烦 ，我这里使用了最简单的卷积加全连接层的模型。

import torch
import torch.nn as nn

# 定义AlexNet模型结构
class AlexNet(nn.Module):
    def __init__(self, num_classes=10):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=11, stride=1, padding=2),
            #nn.ReLU(inplace=True),
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            #nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            #nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            #nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            #nn.ReLU(inplace=True),
        )
        self.classifier = nn.Sequential(
            #nn.Dropout(),
            nn.Linear(256 * 10 * 10, 4096),
            #nn.ReLU(inplace=True),
            #nn.Dropout(),
            nn.Linear(4096, 4096),
            #nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

训练代码（python）

import torch
import torchvision as tv
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
from model import AlexNet

model=AlexNet();

transform = transforms.Compose([transforms.ToTensor(),])

train_dataset=tv.datasets.MNIST(root='./data', train=True, download=True,transform=transform)
test_dataset=tv.datasets.MNIST(root='./data', train=False, download=True,transform=transform)

train_loader = DataLoader(dataset=train_dataset, batch_size=500, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=500, shuffle=False)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

model.to("cuda")
model.train()
num_epochs = 50
for epoch in range(num_epochs):
    for images,labels in train_loader:
        images=images.to("cuda")
        labels=labels.to("cuda")
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
    torch.save(model.state_dict(),"model.pth")

模型导出

import torch

from model import AlexNet

model = AlexNet()

model.load_state_dict(torch.load("model.pth"))

input_example = torch.randn(1, 1, 28, 28)  # 示例输入
torch.onnx.export(model, input_example, 'model.onnx',export_params=True, input_names=['input'], output_names=['output'])

visual studio下C++推理部署

#include <iostream>
#include <assert.h>
#include <onnxruntime_cxx_api.h>
#include <opencv2\opencv.hpp>


int main() {
	const wchar_t*model_path = L"model.onnx";

	Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "rec");
	Ort::SessionOptions session_options;
	Ort::Session session(env, model_path, session_options);
	Ort::AllocatorWithDefaultOptions allocator;

	// 获取输入输出
	std::vector<const char*> input_node_names = { "input" };
	std::vector<const char*> output_node_names = { "output" };

	std::string imgpath = "img.png";//输入你自己的图像路径
	// 这一段代码根据自己的输入而定，我对图像进行了resize和二值化处理
	cv::Mat original_image = cv::imread(imgpath, cv::IMREAD_GRAYSCALE);
	cv::resize(original_image, original_image, cv::Size(28,28));
	cv::threshold(original_image, original_image, 120, 1, cv::THRESH_BINARY);


	for (int y = 0; y < original_image.rows; ++y) {
		for (int x = 0; x < original_image.cols; ++x) {
			uchar pixelValue = original_image.at<uchar>(y, x);
			std::cout << static_cast<int>(pixelValue);
		}
		std::cout <<  std::endl;
	}

	// 确定输入数据维度
	std::vector<int64_t> input_node_dims = { 1,1,28,28 };
	size_t input_tensor_size = 1 * 1 * 28 * 28;

	// 填充数据输入
	std::vector<float> input_tensor_values(input_tensor_size);
	for (int h = 0; h < 28; ++h) {
		for (int w = 0; w < 28; ++w) {
			float pix = original_image.at<uchar>(h, w);
			input_tensor_values[h * 28 + w] = pix;
		}
	}

	auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);

	Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
		memory_info,
		input_tensor_values.data(),
		input_tensor_size,
		input_node_dims.data(),
		input_node_dims.size()
	);

	assert(input_tensor.IsTensor());

	std::vector<Ort::Value> ort_inputs;
	ort_inputs.push_back(std::move(input_tensor));


	// 启动模型预测并获取输出张量
	auto output_tensors = session.Run(
		Ort::RunOptions{ nullptr },
		input_node_names.data(),
		ort_inputs.data(),
		ort_inputs.size(),
		output_node_names.data(),
		1
	);

	Ort::Value& output_tensor = output_tensors[0];
	const float* output_data = output_tensor.GetTensorData<float>();
	float max_index, max_prob=0;

	for (int i = 0; i < 10; i++) {
		if (output_data[i] > max_prob) {
			max_index = i;
			max_prob = output_data[i];
		}
		//std::cout <<i<<" : " << output_data[i] << std::endl;
	}
	std::cout <<"  predict reslut: " << max_index << std::endl;
	
	return 0;
}

感谢观看，如有不足请在评论区指出。