【深度学习】【图像分类】【OnnxRuntime】【C++】ResNet模型部署

提示:博主取舍了很多大佬的博文并亲测有效,分享笔记邀大家共同学习讨论

文章目录

【深度学习】【图像分类】【OnnxRuntime】【C++】ResNet模型部署
前言
Windows平台搭建依赖环境
模型转换--pytorch转onnx
ONNXRuntime推理代码
总结

前言

本期将讲解深度学习图像分类网络ResNet模型的部署，对于该算法的基础知识，可以参考博主【ResNet模型算法Pytorch版本详解】博文。

读者可以通过学习【onnx部署】部署系列学习文章目录的onnxruntime系统学习--C++篇 的内容，系统的学习OnnxRuntime部署不同任务的onnx模型。

Windows平台搭建依赖环境

在【入门基础篇】中详细的介绍了onnxruntime环境的搭建以及ONNXRuntime推理核心流程代码，不再重复赘述。

模型转换--pytorch转onnx

bash 复制代码

import torch
import torchvision as tv
def resnet2onnx():
    # 使用torch提供的预训练权重 1000分类
    model = tv.models.resnet50(pretrained=True)
    model.eval()
    model.cpu()
    dummy_input1 = torch.randn(1, 3, 224, 224)
    torch.onnx.export(model, (dummy_input1), "resnet50.onnx", verbose=True, opset_version=11)
if __name__ == "__main__":
    resnet2onnx()

如下图，torchvision本身提供了不少经典的网络，为了减少教学复杂度，这里博主直接使用了torchvision提供的ResNet网络，并下载和加载了它提供的训练权重。这里可以替换成自己的搭建的ResNet网络以及自己训练的训练权重。

ONNXRuntime推理代码

需要配置imagenet_classes.txt【百度云下载，提取码：rkz7 】文件存储1000类分类标签，假设是用户自定的分类任务，需要根据实际情况作出修改，并将其放置到工程目录下(推荐)。

这里需要将resnet50.onnx放置到工程目录下(推荐)，并且将以下推理代码拷贝到新建的cpp文件中，并执行查看结果。

cpp 复制代码

#include "onnxruntime_cxx_api.h"
#include "cpu_provider_factory.h"
#include <opencv2/opencv.hpp>
#include <fstream>

// 加载标签文件获得分类标签
std::string labels_txt_file = "D:/C++_demo/onnxruntime_onnx/imagenet_classes.txt";
std::vector<std::string> readClassNames();
std::vector<std::string> readClassNames()
{
	std::vector<std::string> classNames;

	std::ifstream fp(labels_txt_file);
	if (!fp.is_open())
	{
		printf("could not open file...\n");
		exit(-1);
	}
	std::string name;
	while (!fp.eof())
	{
		std::getline(fp, name);
		if (name.length())
			classNames.push_back(name);
	}
	fp.close();
	return classNames;
}

int main(int argc, char** argv) {
	// 预测的目标标签数
	std::vector<std::string> labels = readClassNames();

	// 测试图片
	cv::Mat image = cv::imread("D:/C++_demo/onnxruntime_onnx/lion.jpg");
	cv::imshow("输入图", image);

	// 初始化ONNXRuntime环境
	Ort::Env env = Ort::Env(ORT_LOGGING_LEVEL_ERROR, "ResNet-onnx");

	// 设置会话选项
	Ort::SessionOptions session_options;
	// 优化器级别:基本的图优化级别
	session_options.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);
	// 线程数:4
	session_options.SetIntraOpNumThreads(4);
	// 设备使用优先使用GPU而是才是CPU
	std::cout << "onnxruntime inference try to use GPU Device" << std::endl;
	OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0);
	OrtSessionOptionsAppendExecutionProvider_CPU(session_options, 1);

	// onnx训练模型文件
	std::string onnxpath = "D:/C++_demo/onnxruntime_onnx/resnet50.onnx";
	std::wstring modelPath = std::wstring(onnxpath.begin(), onnxpath.end());

	// 加载模型并创建会话
	Ort::Session session_(env, modelPath.c_str(), session_options);

	// 获取模型输入输出信息
	int input_nodes_num = session_.GetInputCount();			// 输入节点输
	int output_nodes_num = session_.GetOutputCount();		// 输出节点数
	std::vector<std::string> input_node_names;				// 输入节点名称
	std::vector<std::string> output_node_names;				// 输出节点名称
	Ort::AllocatorWithDefaultOptions allocator;		
	// 输入图像尺寸
	int input_h = 0;		
	int input_w = 0;

	// 获取模型输入信息
	for (int i = 0; i < input_nodes_num; i++) {
		// 获得输入节点的名称并存储
		auto input_name = session_.GetInputNameAllocated(i, allocator);
		input_node_names.push_back(input_name.get());
		// 显示输入图像的形状
		auto inputShapeInfo = session_.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();
		int ch = inputShapeInfo[1];
		input_h = inputShapeInfo[2];
		input_w = inputShapeInfo[3];
		std::cout << "input format: " << ch << "x" << input_h << "x" << input_w << std::endl;
	}

	// 获取模型输出信息
	int num = 0;
	int nc = 0;
	for (int i = 0; i < output_nodes_num; i++) {
		// 获得输出节点的名称并存储
		auto output_name = session_.GetOutputNameAllocated(i, allocator);
		output_node_names.push_back(output_name.get());
		// 显示输出结果的形状
		auto outShapeInfo = session_.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape();
		num = outShapeInfo[0];
		nc = outShapeInfo[1];
		std::cout << "output format: " << num << "x" << nc << std::endl;
	}

	// 输入数据预处理
	cv::Mat rgb, blob;
	// 默认是BGR需要转化成RGB
	cv::cvtColor(image, rgb, cv::COLOR_BGR2RGB);
	// 对图像尺寸进行缩放
	cv::resize(rgb, blob, cv::Size(input_w, input_h));
	blob.convertTo(blob, CV_32F);
	// 对图像进行标准化处理
	blob = blob / 255.0;	// 归一化
	cv::subtract(blob, cv::Scalar(0.485, 0.456, 0.406), blob);	// 减去均值
	cv::divide(blob, cv::Scalar(0.229, 0.224, 0.225), blob);	//除以方差
	// CHW-->NCHW 维度扩展
	cv::Mat timg = cv::dnn::blobFromImage(blob);
	std::cout << timg.size[0] << "x" << timg.size[1] << "x" << timg.size[2] << "x" << timg.size[3] << std::endl;
	// 占用内存大小,后续计算是总像素*数据类型大小
	size_t tpixels = input_h * input_w * 3;
	std::array<int64_t, 4> input_shape_info{ 1, 3, input_h, input_w };

	// 准备数据输入
	auto allocator_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
	Ort::Value input_tensor_ = Ort::Value::CreateTensor<float>(allocator_info, timg.ptr<float>(), tpixels, input_shape_info.data(), input_shape_info.size());
	
	// 模型输入输出所需数据(名称及其数量),模型只认这种类型的数组
	const std::array<const char*, 1> inputNames = { input_node_names[0].c_str() };
	const std::array<const char*, 1> outNames = { output_node_names[0].c_str() };
	
	// 模型推理
	std::vector<Ort::Value> ort_outputs;
	try {
		ort_outputs = session_.Run(Ort::RunOptions{ nullptr }, inputNames.data(), &input_tensor_, 1, outNames.data(), outNames.size());
	}
	catch (std::exception e) {
		std::cout << e.what() << std::endl;
	}
	// 1x5 获取输出数据并包装成一个cv::Mat对象,为了方便后处理
	const float* pdata = ort_outputs[0].GetTensorMutableData<float>();
	cv::Mat prob(num, nc, CV_32F, (float*)pdata);

	// 后处理推理结果
	cv::Point maxL, minL;		// 用于存储图像分类中的得分最小值索引和最大值索引(坐标)
	double maxv, minv;			// 用于存储图像分类中的得分最小值和最大值
	cv::minMaxLoc(prob, &minv, &maxv, &minL, &maxL); 

	int max_index = maxL.x;		// 获得最大值的索引,只有一行所以列坐标既为索引
	std::cout << "label id: " << max_index << std::endl;
	// 在测试图像上加上预测的分类标签
	cv::putText(image, labels[max_index], cv::Point(50, 50), cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(0, 0, 255), 2, 8);
	cv::imshow("输入图像", image);
	cv::waitKey(0);

	// 释放资源
	session_options.release();
	session_.release();
	return 0;
}

图片正确预测为狮子(lion)：

其实图像分类网络的部署代码基本是一致的，几乎不需要修改，只需要修改传入的图片数据已经训练模型权重即可。

总结

尽可能简单、详细的讲解了C++下onnxruntime环境部署ResNet模型的过程。