pytorch1.3版本转C++ API(libtorch)踩坑记录

前言

最近需要把之前写的一套模型部署到服务器,需要从pytorch->C++,参考官方提供的方案基于libtorch1.3版本。
先列几个我觉得写的比较好的教程:
https://oldpan.me/archives/pytorch-c-libtorch-inference(这篇直接拿一整个工程讲的,很全面。)
https://pytorch.apachecn.org/docs/1.2/advanced/cpp_export.html(1.2版本example官方文档的翻译版本)
1.3版本改动导致的问题:
https://www.jianshu.com/p/186bcdfe9492(必看,因为现在1.2版本的很多教程到1.3跑都有坑)

一、下载安装Libtorch

官网下完解压就好,没踩到坑。
pytorch1.3版本转C++ API(libtorch)踩坑记录_第1张图片

二、模型转换

libtorch提供了trace和script两种方式来对nn.Module进行转换,创建ScriptModule模型,得到的.pt模型能够直接在c++程序中调用。

import torch
from XXX import Net

net = Net()
net.load_state_dict(torch.load("ModelPath", map_location="cpu"))
data = torch.rand(10, 3, 46, 22)
sc = torch.jit.trace(net, data)
sc.save("model.pt")

对于复杂模型,尤其是包含if-else各种逻辑或者非通用库的,可以参考https://blog.csdn.net/xxradon/article/details/86504906,简单sequential模型可以用上面的代码。

三、CMakeList编写

攥写内容大同小异,我参考的https://www.jianshu.com/p/1034b23c703a写的CMakeList,对于语法理解可以参考https://blog.csdn.net/dbzhang800/article/details/6314073。

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(test_pytorch)
set(CMAKE_CXX_STANDARD 11)
set(Torch_DIR /home/XXX/libtorch/share/cmake/Torch) #指定libTorch位置

#include_directories(${OpenCV_INCLUDE_DIRS} /home/zhanyue/tools/libtorch/include /home/zhanyue/tools/libtorch/include/torch/csrc/api/include)
find_package(OpenCV REQUIRED)
find_package(Torch REQUIRED)    # 自动查找libTorch包

四、C++ demo

开始一直跑不通,报错如下

note:   ‘torch::jit::script::Module’ is not derived from ‘const torch::jit::script::slot_iterator_impl<TT>assert(module != nullptr);
error: base operand of ‘->’ has non-pointer type ‘torch::jit::script::Module’
     module->to(device);
           ^
error: base operand of ‘->’ has non-pointer type ‘torch::jit::script::Module’
         at::Tensor output = module->forward(inputs).toTensor();
                                   ^

找了一下午资料,终于在https://www.jianshu.com/p/186bcdfe9492找到原因,pytorch1.3版本Module已经不再是指针了,所以需要去掉断言,然后把->操作换掉,代码如下

#include 
#include "torch/script.h"
#include "torch/torch.h"
using namespace std;

int main() {
    //加载pytorch模型
    torch::jit::script::Module module = torch::jit::load("model.pt");//注意检查路径

    torch::DeviceType device_type; //设置Device类型
    device_type = torch::kCUDA;  //torch::kCUDA  and torch::kCPU
    torch::Device device(device_type, 0);

    //模型转到GPU中去
    module.to(device);
    torch::cuda::is_available(); //判断是否支持GPU加速
    // Create a vector of inputs.
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(torch::ones({1, 3, 46, 224}).to(device));

    //指定执行位置
    // Execute the model and turn its output into a tensor.
    for (int i = 0; i < 100; i++) {
        at::Tensor output = module.forward(inputs).toTensor();
        std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';
    }
}

然后在命令行编译

mkdir build
cd build
cmake .. -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
make

终于编译成功,可以预见今后的坑还会很多,收工。

Scanning dependencies of target test_pytorch
[ 50%] Building CXX object CMakeFiles/test_pytorch.dir/test.cpp.o
[100%] Linking CXX executable test_pytorch
[100%] Built target test_pytorch

你可能感兴趣的:(模型部署与加速)