【深度学习】【图像分类】【OnnxRuntime】【C++】ResNet模型部署
提示:博主取舍了很多大佬的博文并亲测有效,分享笔记邀大家共同学习讨论
前言
本期将讲解深度学习图像分类网络ResNet模型的部署,对于该算法的基础知识,可以参考博主【ResNet模型算法Pytorch版本详解】博文。
读者可以通过学习【onnx部署】部署系列学习文章目录的onnxruntime系统学习–c++篇 的内容,系统的学习OnnxRuntime部署不同任务的onnx模型。
模型转换–pytorch转onnx
首先需要搭建pytorch环境,推荐基础的读者参考博主的博文搭建环境【PyTorch环境搭建】
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网络以及自己训练的训练权重。
Windows平台搭建依赖环境
在【入门基础篇】中详细的介绍了onnxruntime环境的搭建以及ONNXRuntime推理核心流程代码,不再重复赘述。
ONNXRuntime推理代码
需要配置imagenet_classes.txt【百度云下载,提取码:rkz7 】文件存储1000类分类标签,假设是用户自定的分类任务,需要根据实际情况作出修改,并将其放置到工程目录下(推荐)。
这里需要将resnet50.onnx放置到工程目录下(推荐),并且将以下推理代码拷贝到新建的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模型的过程。