C++运行 pytorch(libtorch)模型
前言
应项目需求,需要将pytorch的简单模型移植至ros上,先记录一下C++载入pytorch模型。
安装
很简单,不详细介绍安装方法。
- cmake
- libtorch
- opencv(C++版)
代码
.
├── CMakeLists.txt
├── data
│ ├── llimage.pt
│ └── test_img.jpg
├── src
│ └── cclibtorch.cpp
└── test.sh
执行文件cclibtorch.cpp,其中主要包括加载模型,前向推理,cv::Mat与tensor相互转换等简单操作。
#include <torch/script.h>
#include "opencv2/imgproc.hpp"
#include "opencv2/highgui.hpp"
#include "opencv2/ml.hpp"
#include "opencv2/objdetect.hpp"
#include <iostream>
#include <memory>
int main(int argc, const char* argv[]){
if(argc != 3){
std::cerr << "usage: main <path-to-exported-script-module and image>\n";
return -1;
}
torch::Device device(torch::kCUDA);
torch::jit::script::Module module = torch::jit::load(argv[1]); // load model
module.to(device);
module.eval();
cv::Mat input_img = cv::imread(argv[2]);
if(!input_img.data){
std::cerr << "Problem loading image!!!" << std::endl;
return 0;
}
cv::imshow("input_img",input_img);
cv::cvtColor(input_img, input_img, cv::COLOR_BGR2RGB);
int img_h = input_img.rows;
int img_w = input_img.cols;
input_img.convertTo(input_img, CV_32FC3, 1.0f / 255.0f);
auto input_tensor = torch::from_blob(input_img.data, {1, img_h, img_w, 3});
input_tensor = input_tensor.permute({0, 3, 1, 2});
input_tensor = input_tensor.to(device); //send tensor to gpu
//return multi elements is toTuple(),and single element is toTensor();
auto output_result = module.forward({input_tensor}).toTuple();
int elements_num = output_result->elements().size();
for (size_t i = 0; i < elements_num; i++){
torch::Tensor output_tensor = output_result->elements()[i].toTensor();
output_tensor = output_tensor.squeeze().detach().permute({1, 2, 0}).contiguous();
output_tensor = output_tensor.mul(255).clamp(0, 255).to(torch::kU8);
output_tensor = output_tensor.to(torch::kCPU);
cv::Mat cv_result_img(img_h, img_w, CV_8UC3);
std::memcpy((void *) cv_result_img.data, output_tensor.data_ptr(), sizeof(torch::kU8) * output_tensor.numel());
cv::imshow("img_" + std::to_string(i),cv_result_img);
}
std::cerr << "model inference is ok !!!" << std::endl;
cv::waitKey(0);
return 1;
}
CMakeLists.txt文件
cmake_minimum_required(VERSION 3.0)
project(cclibtorch_test)
find_package(OpenCV REQUIRED)
include_directories(${OpenCV_INCLUDE_DIRS})
find_package(PythonInterp REQUIRED) #加载libtorch所安装的路径
set(CMAKE_PREFIX_PATH /home/simba/WorkCode/Env/libtorch)
include_directories(/home/simba/WorkCode/Env/libtorch/include)
find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
MESSAGE(STATUS "Project: ${PROJECT_NAME}")
MESSAGE(STATUS "OpenCV library status:")
MESSAGE(STATUS " version: ${OpenCV_VERSION}")
MESSAGE(STATUS " libraries: ${OpenCV_LIBS}")
MESSAGE(STATUS " include path: ${OpenCV_INCLUDE_DIRS}")
add_executable(cclibtorch_test src/cclibtorch.cpp)
target_compile_features(cclibtorch_test PUBLIC cxx_range_for)
target_link_libraries(cclibtorch_test ${TORCH_LIBRARIES} ${OpenCV_LIBS})
set_property(TARGET cclibtorch_test PROPERTY CXX_STANDARD 14)
终端执行操作
mkdir build && cd build
cmake .. && make -j4
cd ..
./build/cclibtorch_test ./data/llimage.pt ./data/test_img.jpg
小程序链接
小程序编译链接,提取码: 567u
参考文章
感谢各位大佬的分享
下雨天晴77—C++ opencv矩阵和pytorch tensor的互相转换
峰峰的猫—pytorch的tensor,Image,numpy和opencv四种格式的相互转换
农夫山泉2号—cv mat转tensor(b,c,h,w)