2023年11月29日发(作者:)

C++PytorchOnnxRuntime使⽤⽅法(附代码)

PytorchC++OnnxRuntime使⽤⽅法

写在前⾯

如何使⽤?

改⼀下代码中的输⼊图⽚路径、模型路径、输出图⽚路径即可~

下⾯对⼀些需要注意的细节做⼀下解释

Pytoch onnxruntime使⽤⽅法

⼤概步骤就是,先保存训练好的模型,然后将该模型转为onnx模型,使⽤onnxruntime运⾏模型即可。具体如何使⽤可以去看和我的

GitHub代码,备注写的满详细的。下⾯对⼀些需要注意的地⽅做⼀些讲解。

1、安装onnxrutime

pip install onnxruntime-gpu

2、 导出onnx模型,详见export_⽂件

需要注意的是最终导出onnx模型时,你的batch_size设为多少,导出的onnx模型就只能接收batch_size⼤⼩的图⽚,⽐如说这⾥

batch_size为1,表⽰你之后在⽤onnx模型时模型的输⼊只能是[1,3,256,256],不能是[2,3,256,256]。

3、验证导出的onnx模型有没有问题,详见verify_

如果onnx模型输出是你想要的⼤⼩那应该就没问题。

4、运⾏onnx模型,详见run_

这个没什么好说的,去官⽹看看例⼦就知道了,核⼼就是这⼏句代码。

ort_session = onnxruntime.InferenceSession("pix2pixHD_Cartoon_")

ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(label.unsqueeze(0)).astype(np.float32)}

ort_outs = ort_session.run(None, ort_inputs)

C++ onnxruntime使⽤⽅法

重点是C++如何使⽤onnxruntime,感觉给的例⼦并不是很好。其实个⼈感觉最难的还是如何进⾏预处理和后处理,因为C++不像Python

那样可以灵活使⽤矩阵,如果⽤for循环⼀个像素⼀个像素的进⾏处理,那速度就太慢了~,所以我⽤到了opencv、libtorch这两个库来对

数据进⾏处理。具体如何使⽤可以看我放在GitHub上的代码。下⾯对⼀些个⼈感觉⽐较重要的地⽅进⾏解释

1、安装Onnxruntime-Gpu

这⾥建议安装1.7.1 因为安装最新的版本会报错~

2、Libtorch的安装

要注意的是链接器输⼊添加:

asmjit.lib

c10.lib

caffe2_detectron_ops.lib

caffe2_module_test_dynamic.lib

clog.lib

cpuinfo.lib

dnnl.lib

fbgemm.lib

libprotobuf.lib

libprotobuf-lite.lib

libprotoc.lib

mkldnn.lib

torch.lib

torch_cpu.lib

因为不同版本的libtorch不⼀样,可能不全需要,保险起见可以全加上

3、将Opencv的Mat对象转为Libtorch的tensor对象这样就⽅便使⽤API对图像进⾏归⼀化处理

Mat转tensor的代码,需要注意的是torch::kByte要对应cv::Mat img的数据类型

torch::Tensor img_tensor = torch::from_blob(img.data, { img.rows, img.cols, 3 }, torch::kByte);

4、归⼀化

转为tensor过后进⾏归⼀化操作就很简单了⽽且速度还快,嘿嘿~

std::vector<double> mean = { 0.5,0.5,0.5 };

std::vector<double> std = { 0.5,0.5,0.5 };

tensor = torch::data::transforms::Normalize<>(mean, std)(tensor);

5、构造onnx模型的输⼊

std::vector inputVec

个⼈感觉这是⼀个难点,之前⽤for循环⼀个⼀个遍历图像放进去的,这样太慢了。后来想到可以⽤内存拷贝的⽅法。实现的思路就是将已

经预处理好的图像分成RGB三个通道然后reshape成⼀维向量装⼊inputVec⾥。 为什么要分为三个通道装⼊呢,直接将⼀张图像reshape

后装⼊inputVec它不⾹吗?这是因为inputVec向量的输⼊顺序要求是先R通道的所有像素再G通道的所有像素最后再G通道的所有像素,即

要求顺序为RRR->GGG->BBB。如果直接将图像reshape为⼀维向量则inputVec的顺序为RGB->RGB->RGB。。。。最后结果就不对啦~

核⼼代码如下:

cv::Mat resultImg(h, w, CV_32FC3);

std::memcpy((void*)resultImg.data, label.data_ptr(), sizeof(float) * label.numel());

//

通道分离

std::vector<cv::Mat> channels;

cv::split(resultImg, channels);

cv::Mat blue, green, red;

blue = channels.at(0);

green = channels.at(1);

red = channels.at(2);

//

拉平

std::vector<float> inputVec_red = (std::vector<float>)(blue.reshape(1, 1));

std::vector<float> inputVec_green = (std::vector<float>)(green.reshape(1, 1));

std::vector<float> inputVec_blue = (std::vector<float>)(red.reshape(1, 1));

inputVec.insert(inputVec.end(), inputVec_red.begin(), inputVec_red.end());

//

装⼊向量

inputVec.insert(inputVec.end(), inputVec_green.begin(), inputVec_green.end());

inputVec.insert(inputVec.end(), inputVec_blue.begin(), inputVec_blue.end());

这⾥第⼆⾏我是将之前预处理好的tensor转为Mat后再分离通道的,要注意是的是std::memcpy ⾏中sizeof()中内容,需要修改成c++中内

建的数据类型,如果使⽤torch::kF32或者其他浮点型,会出现数据复制缺失的情况,所以使⽤的是float。

6、将构造好的输⼊装⼊模型得到输出

Ort::Env env = Ort::Env{ ORT_LOGGING_LEVEL_ERROR, "Default" };

Ort::SessionOptions opt;

opt.AppendExecutionProvider_CUDA(cuda_options);

Ort::Session session(env, L".onnx模型的地址", opt);

std::vector<int64_t> inputSize = { batchSize, channels, height, width };

rsize_t inputSizeCount = batchSize * channels * height * width;

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

std::clock_t startTime, endTime;

Ort::Value inputTensor = Ort::Value::CreateTensor<float>(memoryInfo, inputVec.data(), inputSizeCount, inputSize.data(), inputSize.size());

std::vector<Ort::Value> outputTensors = session.Run(Ort::RunOptions{ nullptr }, inputNames.data(), &inputTensor, inputNames.size(), outputNames.data(

), outputNames.size());

float* output = outputTensors[0].GetTensorMutableData<float>();

可以看到输出是⼀个指针

float* output

7、将输出做后处理后转为图像

同理输出的float* 顺序为RRR->GGG->BBB所以如果直接resahpe那是不⾏滴。所以先构造tensor把三个通道分别提出来,再将三通道进

⾏合并。

torch::Tensor result_r = torch::from_blob(output, { height,width,1 });

torch::Tensor result_g = torch::from_blob(&output[height * width - 1], { height,width,1 });

torch::Tensor result_b = torch::from_blob(&output[height * width * 2 - 1], { height,width,1 });

torch::Tensor result = torch::cat({ result_r, result_g, result_b }, 2);

然后使⽤API直接进⾏后处理,这样就很⽅便。记得转为类型因为后⾯就要涉及到内存拷贝了

torch::kU8

result = result.add(1).div(2).mul(255);

result = result.clamp(0, 255);

result = result.to(torch::kU8);

最后再将tensor转为Mat进⾏保存

cv::Mat resultImg(height, width, CV_8UC3);

std::memcpy((void*)resultImg.data, result.data_ptr(), sizeof(torch::kU8) * result.numel());

cv::cvtColor(resultImg, resultImg, cv::COLOR_RGB2BGR);

这⾥我看⽹上代码有将tensor先permute为HWC后再转为Mat,经过我实验发现其实这⼀步是多余的,因为permute⽅法并不会改变

tensor数据在内存中的排列⽅式,对于拷贝内存来说是多余的⼀步。

写在后⾯

⽬前我能想到的地⽅就这些了,如果错误的地⽅请多多包含啦~欢迎⼤家提问!