使用C++調用pytorch模型(Linux)

前言

模型轉換思路通常為:

• Pytorch -> ONNX -> TensorRT
• Pytorch -> ONNX -> TVM
• Pytorch -> 轉換工具 -> caffe
• Pytorch -> torchscript(C++版本Torch)

---

我的模型是使用Pytorch1.0訓練的,第三種方法應該是還不支持,沒有對應層名字, 放棄. (以下是用方法3生成的網絡結構圖, 其中部分層名字和工具對應不上).

因此本文使用第4中方法,詳細步驟分兩步, 具體如下(目前資料少,坑很多)

---

1. pytorch模型轉化為libtorch的torchscript模型 (.pth -> .pt)

首先, 在python中, 把模型轉化成.pt文件

Pytorch官方提供的C++API名為libtorch，詳細查看:
- LIBRARY API
- USING THE PYTORCH C++ FRONTEND

import torch

# An instance of your model.
from my_infer import BaseLine

model = BaseLine().model.cpu().eval()

# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 256 , 128)

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("demo/model.pt")

2. 使用libtorch調用torchscript模型

此處有一個大坑, opencv和torch可以單獨使用, 但如果鏈接libtorch庫以后, cv::imread提示未定義的應用. 所以使用了opencv2的圖片讀取方式, 然后再轉成cv::Mat.

---

更新時間:2019/05/24
在更換libtorch版本后, cv:imread不再報錯, 具體原因說不上來, 應該是之前的版本鏈接庫時候出現矛盾什么的...

---

#include <iostream>                                                                                                
#include "torch/script.h"
#include "torch/torch.h"
#include "opencv2/core.hpp"
#include "opencv2/imgproc.hpp"
#include "opencv2/highgui.hpp"
#include <vector>

int main()
{
    //加載pytorch模型
    std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("/home/zhuoshi/ZSZT/Geoffrey/opencvTest/m
    assert(module != nullptr);

    // 創建一個Tensor
    //std::vector<torch::jit::IValue> inputs;
    //inputs.emplace_back(torch::ones({1, 3, 256, 128}));
    //測試前向
    //at::Tensor output = module->forward(inputs).toTensor();
    //std::cout << output;

    // 轉換為int8類型
    //vector<int16_t> feature(2048);
    //for (int i = 0;i<128;i++)
    //{
    // 轉化成Float
    //int temp = output[0][i].item().toInt();
    //    if (temp != 0){
    //        temp = 1;
    //    }
    //    feature[i] = temp;
    //}
    //std::cout << feature;

    //讀取圖片
    IplImage* pmg = cvLoadImage("/home/zhuoshi/ZSZT/Geoffrey/opencvTest/test.jpg");
    cv::Mat image(pmg, true);
    //cv::Mat imageRGB = cv::cvtColor(image, imageRGB, cv::COLOR_BGR2RGB);
    cv::cvtColor(image, image, CV_BGR2RGB);

    //IplImage轉換成Tensor
    cv::Mat img_float;
    image.convertTo(img_float, CV_32F, 1.0 / 255);
    cv::resize(img_float, img_float, cv::Size(256, 128));
    torch::Tensor tensor_image = torch::from_blob(img_float.data, {1, 3, 256, 128}, torch::kFloat32);

    //前向
    std::vector<torch::jit::IValue> input;
    input.emplace_back(tensor_image);
    at::Tensor output_image = module->forward(input).toTensor();
    //std::cout << output_image;

    //Tensor 轉 array
    std::vector<float> feature(2048);
    for (int i=0; i<2048; i++){
    //    feature[i] = output_image[i]
        std::cout << output_image[0][i].item().toFloat();
    }
    return 0;
}  

對應的CMakeLists.txt內容:

cmake_minimum_required(VERSION 2.8)                                                                                

project(opencv_example_project)
SET(CMAKE_C_COMPILER g++)
add_definitions(--std=c++11)

# 指定libTorch位置
set(Torch_DIR /home/zhuoshi/ZSZT/Geoffrey/opencvTest/libtorch/share/cmake/Torch)
find_package(Torch REQUIRED)

find_package(OpenCV REQUIRED)

message(STATUS "OpenCV library status:")
message(STATUS "    version: ${OpenCV_VERSION}")
message(STATUS "    libraries: ${OpenCV_LIBS}")
message(STATUS "    include path: ${OpenCV_INCLUDE_DIRS}")
message(STATUS "    torch lib : ${TORCH_LIBRARIES} ")

include_directories(${OpenCV_INCLUDE_DIRS}
                    /home/zhuoshi/ZSZT/Geoffrey/opencvTest/libtorch/include
                    /home/zhuoshi/ZSZT/Geoffrey/opencvTest/libtorch/include/torch/csrc/api/include/
                    )

add_executable(main main.cpp)
     
# Link your application with OpenCV libraries
target_link_libraries(main ${OpenCV_LIBS} ${TORCH_LIBRARIES} )

運行結果如圖:

---

更新時間: 2019/05/25, 更換libtorch版本后, cv::read可用, 這是新版本

#include <iostream>
#include "torch/script.h"
#include "torch/torch.h"
#include "opencv2/core.hpp"
#include "opencv2/imgproc.hpp"
#include "opencv2/highgui.hpp"
#include "opencv2/imgcodecs.hpp"
#include <vector>

int main()
{
    /* 配置參數 */
    std::vector <float> mean_ = {0.485, 0.456, 0.406};
    std::vector <float> std_ = {0.229, 0.224, 0.225};
    char path[] = "../test.jpg";

    // 讀取圖片
    cv::Mat image = cv::imread(path);
    if (image.empty())
        fprintf(stderr, "Can not load image\n");

    // 轉換通道,
    cv::cvtColor(image, image, CV_BGR2RGB);
    cv::Mat img_float;
    image.convertTo(img_float, CV_32F, 1.0 / 255);

    // resize, 測試一個點數據
    cv::resize(img_float, img_float, cv::Size(256, 128));
    //std::cout << img_float.at<cv::Vec3f>(256, 128)[1] << std::endl;

    // 轉換成tensor
    auto img_tensor = torch::from_blob(img_float.data, {1, 3, 256, 128}, torch::kFloat32);
    //img_tensor = img_tensor.permute({0,3,1,2});
    // tensor標准化
    for (int i = 0; i < 3; i++) {
        img_tensor[0][0] = img_tensor[0][0].sub_(mean_[i]).div_(std_[i]);
    }

    // 構造input
    //auto img_var = torch::autograd::make_variable(img_tensor, false); //tensor->variable會報錯
    std::vector<torch::jit::IValue> inputs;
    inputs.emplace_back(img_tensor); //向容器中加入新的元素, 右值引用

    //加載pytorch模型
    std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("../model/model_int.pt");
    assert(module != nullptr);

    //前向
    at::Tensor output_image = module->forward(inputs).toTensor();
    std::cout << output_image;

    return 0;
}

cv::Mat convertTo3Channels(cv::Mat binImg)
{
    cv::Mat three_channel = cv::Mat::zeros(binImg.rows, binImg.cols, CV_8UC3);
    std::vector<cv::Mat> channels;
    for (int i=0;i<3;i++)
    {
        channels.push_back(binImg);
    }
    merge(channels, three_channel);
    return three_channel;
}

對應CMakelist.txt文件:

cmake_minimum_required(VERSION 2.8)

# Define project name
project(opencv_example_project)

SET(CMAKE_C_COMPILER g++)
add_definitions(--std=c++11)

# 指定libTorch位置
set(Torch_DIR /home/geoffrey/CLionProjects/opencvTest/libtorch/share/cmake/Torch)
find_package(Torch REQUIRED)

message(STATUS "Torch library status:")
message(STATUS "    version: ${TORCH_VERSION}")
message(STATUS "    libraries: ${TORCH_LIBS}")
message(STATUS "    include path: ${TORCH_INCLUDE_DIRS}")
message(STATUS "    torch lib : ${TORCH_LIBRARIES} ")

# 指定OpenCV位置
#set(OpenCV_DIR /run/media/geoffrey/Timbersaw/Backup/other_package/opencv-4.0.0/build)
# set(OpenCV_DIR /opt/opencv2)
find_package(OpenCV  REQUIRED)
message(STATUS "OpenCV library status:")
message(STATUS "    version: ${OpenCV_VERSION}")
message(STATUS "    libraries: ${OpenCV_LIBS}")
message(STATUS "    include path: ${OpenCV_INCLUDE_DIRS}")
message(STATUS "    opencv lib : ${OpenCV_LIBRARIES} ")

# 包含頭文件include
include_directories(${OpenCV_INCLUDE_DIRS} ${TORCH_INCLUDE_DIRS})

# 生成的目標文件（可執行文件）
add_executable(main main.cpp)

# 置需要的庫文件lib
# set(OpenCV_LIBS opencv_core  opencv_highgui opencv_imgcodecs opencv_imgproc)
target_link_libraries(main  ${OpenCV_LIBS} ${TORCH_LIBRARIES}) #

---

---

參考資料

1. TorchDemo
2. 利用Pytorch的C++前端(libtorch)讀取預訓練權重並進行預測
3. C++部署pytorch模型（二）————使用libtorch調用torchscripts模型

---