将PyTorch模型部署到C++上(Ubuntu)

  1. 为什么要这么做?

PyTorch作为一个开源的Python机器学习库正受到越来越广泛的关注和应用。然而,作为一门语言,Python并不是在任何场景下都适用的。在生产以及部署到自动驾驶车辆上等场景中,C++常常是更好的选择。因此需要将PyTorch模型部署到C++上。以下是实现的步骤。

  1. 将PyTorch模型转化成Torch Script

Torch Script的作用是用TorchScript编写的任何代码都可以从Python进程中保存并加载到没有Python依赖关系的进程中。有两种方法可以实现这一步骤。

(1)tracing

在这种机制中,通过使用示例输入对模型进行分析,并记录通过模型的这些输入流,来捕获模型的结构。这适用于控制流使用有限的模型。具体做法为,将示例输入到模型中,并将其作为一个例子输入到torch.jit.trace函数中。这将产生一个torch.jit.ScriptModule对象,并在模块的forward方法中嵌入模型跟踪,具体代码如下:

import torch
import torchvision

# An instance of your model.
model = torchvision.models.resnet18()

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

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)

(2)使用注释来转化

在某些情况下,比如模型采用了特定形式的控制流,可能希望直接用Torch Script编写模型,并相应地对模型进行注释。比如以下这个模型:

import torch

class MyModule(torch.nn.Module):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    def forward(self, input):
        if input.sum() > 0:
          output = self.weight.mv(input)
        else:
          output = self.weight + input
        return output

由于该模块的forward方法使用依赖于输入的控制流,因此不适合tracing。相反,我们可以将其转换为ScriptModule。为了将模块转换为ScriptModule,需要使用torch.jit.script编译模块,如下所示:

class MyModule(torch.nn.Module):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    def forward(self, input):
        if input.sum() > 0:
          output = self.weight.mv(input)
        else:
          output = self.weight + input
        return output

my_module = MyModule(10,20)
sm = torch.jit.script(my_module)
  1. 将Script Module输出为文件

(1) tracing方法:

traced_script_module.save("traced_resnet_model.pt")

(2)script方法:

sm.save("my_module_model.pt")

到这步我们就已经结束了在Python上的操作,要开始转战C++了。

  1. 在C++中载入Script Module

要用c++加载序列化的PyTorch模型,应用程序必须依赖于LibTorch。LibTorch发行版包含了一组共享库、头文件和CMake配置文件。虽然CMake不是必须的,但它是推荐的方法,并且在将来会得到很好的支持。在本教程中,我们将使用CMake和LibTorch构建一个c++应用程序,它只是用于加载和执行一个序列化的PyTorch模型。

LibTorch的下载地址在这里,根据自己的需求和电脑的配置下载,本文对应的版本是1.8.1CPU版,其他版本的也许本文的方法不适用,请注意!

首先,需要编写代码来载入模块。

#include <torch/script.h> // One-stop header.

#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
  if (argc != 2) {
    std::cerr << "usage: example-app <path-to-exported-script-module>\n";
    return -1;
  }


  torch::jit::script::Module module;
  try {
    // Deserialize the ScriptModule from a file using torch::jit::load().
    module = torch::jit::load(argv[1]);
  }
  catch (const c10::Error& e) {
    std::cerr << "error loading the model\n";
    return -1;
  }

  std::cout << "ok\n";
}

接着,创建一个CMakeLists.txt, 内容如下:

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)//写项目名

//set(Torch_DIR /home/nio/libtorch-1.8.1/libtorch/share/cmake/Torch)
find_package(Torch REQUIRED)

add_executable(example-app example-app.cpp)//文件夹名,文件名
target_link_libraries(example-app "${TORCH_LIBRARIES}")//文件夹名
set_property(TARGET example-app PROPERTY CXX_STANDARD 14)//文件夹名

这里需要注意,按官网教程,是不用set行的,在很多别人的教程里也不用,但我在实验的时候,没有set行设置Torch地址,后续的操作会报错,提示找不到Torch,所以看情况自己加。

两个文件创建好之后,放在一个文件夹下,如:

example-app/
  CMakeLists.txt
  example-app.cpp

在该文件夹下,打开终端,输入以下指令:

mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..//背后两个小点别忘了!!!

一开始我没有设置Torch地址,得到的结果为:

-- The C compiler identification is GNU 7.5.0
-- The CXX compiler identification is GNU 7.5.0
-- Check for working C compiler: /usr/bin/cc
-- Check for working C compiler: /usr/bin/cc -- works
-- Detecting C compiler ABI info
-- Detecting C compiler ABI info - done
-- Detecting C compile features
-- Detecting C compile features - done
-- Check for working CXX compiler: /usr/bin/c++
-- Check for working CXX compiler: /usr/bin/c++ -- works
-- Detecting CXX compiler ABI info
-- Detecting CXX compiler ABI info - done
-- Detecting CXX compile features
-- Detecting CXX compile features - done
CMake Error at CMakeLists.txt:4 (find_package):
  By not providing "FindTorch.cmake" in CMAKE_MODULE_PATH this project has
  asked CMake to find a package configuration file provided by "Torch", but
  CMake did not find one.

  Could not find a package configuration file provided by "Torch" with any of
  the following names:

    TorchConfig.cmake
    torch-config.cmake

  Add the installation prefix of "Torch" to CMAKE_PREFIX_PATH or set
  "Torch_DIR" to a directory containing one of the above files.  If "Torch"
  provides a separate development package or SDK, be sure it has been
  installed.


-- Configuring incomplete, errors occurred!
See also "/home/nio/C++FILE/code/vectornet/build/CMakeFiles/CMakeOutput.log".

设置了之后:

-- The C compiler identification is GNU 7.5.0
-- The CXX compiler identification is GNU 7.5.0
-- Check for working C compiler: /usr/bin/cc
-- Check for working C compiler: /usr/bin/cc -- works
-- Detecting C compiler ABI info
-- Detecting C compiler ABI info - done
-- Detecting C compile features
-- Detecting C compile features - done
-- Check for working CXX compiler: /usr/bin/c++
-- Check for working CXX compiler: /usr/bin/c++ -- works
-- Detecting CXX compiler ABI info
-- Detecting CXX compiler ABI info - done
-- Detecting CXX compile features
-- Detecting CXX compile features - done
-- Looking for pthread.h
-- Looking for pthread.h - found
-- Looking for pthread_create
-- Looking for pthread_create - not found
-- Looking for pthread_create in pthreads
-- Looking for pthread_create in pthreads - not found
-- Looking for pthread_create in pthread
-- Looking for pthread_create in pthread - found
-- Found Threads: TRUE  
-- Found Torch: /home/nio/libtorch-1.8.1/libtorch/lib/libtorch.so  
-- Configuring done
-- Generating done
-- Build files have been written to: /home/nio/C++FILE/code/vectornet/build

成功之后,在build文件夹下用make编译:

(base) nio@LT5CG052BHT2:~/C++FILE/code/vectornet/build$ make
Scanning dependencies of target vectornet
[ 50%] Building CXX object CMakeFiles/vectornet.dir/pytoc.cpp.o
[100%] Linking CXX executable vectornet
[100%] Built target vectornet

运行pt文件

(base) nio@LT5CG052BHT2:~/C++FILE/code/vectornet/build$ ./vectornet ../traced_vectornet_model.pt 
ok

最后,执行输入并验证输出。在之前的cpp文件中加入:

// Create a vector of inputs.
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}));

// Execute the model and turn its output into a tensor.
at::Tensor output = module.forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

再make和执行pt文件即可得到结果。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值