2023年11月29日发(作者:)
C++和PytorchOnnxRuntime使⽤⽅法(附代码)
Pytorch和C++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
个⼈感觉这是⼀个难点,之前⽤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数据在内存中的排列⽅式,对于拷贝内存来说是多余的⼀步。
写在后⾯
⽬前我能想到的地⽅就这些了,如果错误的地⽅请多多包含啦~欢迎⼤家提问!


发布评论