[小程序系列] C++ 运行pytorch(libtorch)模型

C++运行 pytorch(libtorch)模型

前言

应项目需求，需要将pytorch的简单模型移植至ros上，先记录一下C++载入pytorch模型。

安装

很简单，不详细介绍安装方法。

1. cmake
2. libtorch
3. opencv(C++版)

代码

.
├── CMakeLists.txt
├── data
│   ├── llimage.pt
│   └── test_img.jpg
├── src
│   └── cclibtorch.cpp
└── test.sh

执行文件cclibtorch.cpp，其中主要包括加载模型，前向推理，cv::Mat与tensor相互转换等简单操作。

#include 
#include "opencv2/imgproc.hpp"
#include "opencv2/highgui.hpp"
#include "opencv2/ml.hpp"
#include "opencv2/objdetect.hpp"
#include 
#include 

int main(int argc, const char* argv[]){
    if(argc != 3){
        std::cerr << "usage: main \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)