基于C++的PyTorch模型部署-技术圈

点击上方“小白学视觉”，选择加"星标"或“置顶”

重磅干货，第一时间送达

引言

PyTorch作为一款端到端的深度学习框架，在1.0版本之后已具备较好的生产环境部署条件。除了在web端撰写REST API进行部署之外(参考)，软件端的部署也有广泛需求。尤其是最近发布的1.5版本，提供了更为稳定的C++前端API。

工业界与学术界最大的区别在于工业界的模型需要落地部署，学界更多的是关心模型的精度要求，而不太在意模型的部署性能。一般来说，我们用深度学习框架训练出一个模型之后，使用Python就足以实现一个简单的推理演示了。但在生产环境下，Python的可移植性和速度性能远不如C++。所以对于深度学习算法工程师而言，Python通常用来做idea的快速实现以及模型训练，而用C++作为模型的生产工具。目前PyTorch能够完美的将二者结合在一起。实现PyTorch模型部署的核心技术组件就是TorchScript和libtorch。

所以基于PyTorch的深度学习算法工程化流程大体如下图所示：

TorchScript

TorchScript可以视为PyTorch模型的一种中间表示，TorchScript表示的PyTorch模型可以直接在C++中进行读取。PyTorch在1.0版本之后都可以使用TorchScript的方式来构建序列化的模型。TorchScript提供了Tracing和Script两种应用方式。

Tracing应用示例如下：

class MyModel(torch.nn.Module):    def __init__(self):        super(MyModel, self).__init__()        self.linear = torch.nn.Linear(4, 4)
    def forward(self, x, h):        new_h = torch.tanh(self.linear(x) + h)        return new_h, new_h
# 创建模型实例 my_model = MyModel()# 输入示例x, h = torch.rand(3, 4), torch.rand(3, 4)# torch.jit.trace方法对模型构建TorchScripttraced_model = torch.jit.trace(my_model, (x, h))# 保存转换后的模型traced_model.save('model.pt')

在这段代码中，我们先是定义了一个简单模型并创建模型实例，然后给定输入示例，Tracing方法最关键的一步在于使用torch.jit.trace方法对模型进行TorchScript转化。我们可以获得转化后的traced_model对象获得其计算图属性和代码属性。计算图属性：

print(traced_model.graph)

graph(%self.1 : __torch__.torch.nn.modules.module.___torch_mangle_1.Module,      %input : Float(3, 4),      %h : Float(3, 4)):  %19 : __torch__.torch.nn.modules.module.Module = prim::GetAttr[name="linear"](%self.1)  %21 : Tensor = prim::CallMethod[name="forward"](%19, %input)  %12 : int = prim::Constant[value=1]() # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0  %13 : Float(3, 4) = aten::add(%21, %h, %12) # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0  %14 : Float(3, 4) = aten::tanh(%13) # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0  %15 : (Float(3, 4), Float(3, 4)) = prim::TupleConstruct(%14, %14)  return (%15)

代码属性：

print(traced_cell.code)


def forward(self,    input: Tensor,    h: Tensor) -> Tuple[Tensor, Tensor]:  _0 = torch.add((self.linear).forward(input, ), h, alpha=1)  _1 = torch.tanh(_0)  return (_1, _1)

这样我们就可以将整个模型都保存到硬盘上了，并且经过这种方式保存下来的模型可以加载到其他其他语言环境中。

TorchScript的另一种实现方式是Script的方式，可以算是对Tracing方式的一种补充。当模型代码中含有if或者for-loop等控制流程序时，使用Tracing方式是无效的，这时候可以采用Script方式来进行实现TorchScript。实现方法跟Tracing差异不大，关键在于把jit.tracing换成jit.script方法，示例如下。

scripted_model = torch.jit.script(MyModel)scripted_model.save('model.pt')

除了Tracing和Script之外，我们也可以混合使用这两种方式，这里不做详述。总之，TorchScript为我们提供了一种表示形式，可以对代码进行编译器优化以提供更有效的执行。

libtorch

在Python环境下对训练好的模型进行转换之后，我们需要C++环境下的PyTorch来读取模型并进行编译部署。这种C++环境下的PyTorch就是libtorch。因为libtorch通常用来作为PyTorch模型的C++接口，libtorch也称之为PyTorch的C++前端。

我们可以直接从PyTorch官网下载已经编译好的libtorch安装包，当然也可以下载源码自行进行编译。这里需要注意的是，安装的libtorch版本要与Python环境下的PyTorch版本一致。

安装好libtorch后可简单测试下是否正常。比如我们用TorchScript转换一个预训练模型，示例如下：

import torchimport torchvision.models as modelsvgg16 = models.vgg16()example = torch.rand(1, 3, 224, 224).cuda() model = model.eval()traced_script_module = torch.jit.trace(model, example)output = traced_script_module(torch.ones(1,3,224,224).cuda())traced_script_module.save('vgg16-trace.pt')print(output)

输出为：

tensor([[ -0.8301, -35.6095, 12.4716]], device='cuda:0',        grad_fn=<AddBackward0>)

然后切换到C++环境，编写CmakeLists文件如下：

cmake_minimum_required(VERSION 3.0.0 FATAL_ERROR)project(libtorch_test)find_package(Torch REQUIRED)message(STATUS "Pytorch status:")message(STATUS "libraries: ${TORCH_LIBRARIES}")add_executable(libtorch_test test.cpp)target_link_libraries(libtorch_test "${TORCH_LIBRARIES}")set_property(TARGET libtorch_test PROPERTY CXX_STANDARD 11)

继续编写test.cpp代码如下：

#include "torch/script.h"#include "torch/torch.h"#include #include using namespace std;
int main(int argc, const char* argv[]){    if (argc != 2) {        std::cerr << "usage: example-app \n";        return -1;    }
    // 读取TorchScript转化后的模型    torch::jit::script::Module module;    try {        module = torch::jit::load(argv[1]);    }
    catch (const c10::Error& e) {        std::cerr << "error loading the model\n";        return -1;    }
    module->to(at::kCUDA);    assert(module != nullptr);    std::cout << "ok\n";
    // 构建示例输入    std::vector inputs;    inputs.push_back(torch::ones({1, 3, 224, 224}).to(at::kCUDA));
    // 执行模型推理并输出tensor    at::Tensor output = module->forward(inputs).toTensor();    std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';}

编译test.cpp并执行，输出如下。对比Python环境下的的运行结果，可以发现基本是一致的，这也说明当前环境下libtorch安装没有问题。

ok-0.8297, -35.6048, 12.4823[Variable[CUDAFloatType]{1,3}]

完整部署流程

通过前面对TorchScript和libtorch的描述，其实我们已经基本将PyTorch的C++部署已经基本讲到了，这里我们再来完整的理一下整个流程。基于C++的PyTorch模型部署流程如下。

第一步：

通过torch.jit.trace方法将PyTorch模型转换为TorchScript，示例如下：

import torchfrom torchvision.models import resnet18model =resnet18()example = torch.rand(1, 3, 224, 224)tracing.traced_script_module = torch.jit.trace(model, example)

第二步：

将TorchScript序列化为.pt模型文件。

traced_script_module.save("traced_resnet_model.pt")

第三步：

在C++中导入序列化之后的TorchScript模型，为此我们需要分别编写包含调用程序的cpp文件、配置和编译用的CMakeLists.txt文件。CMakeLists.txt文件示例内容如下：

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)project(custom_ops)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)

包含模型调用程序的example-app.cpp示例编码如下：

#include  // torch头文件.#include #include 
int main(int argc, const char* argv[]) {  if (argc != 2) {    std::cerr << "usage: example-app \n";    return -1;  }
  torch::jit::script::Module module;  try {    // 反序列化：导入TorchScript模型    module = torch::jit::load(argv[1]);  }
  catch (const c10::Error& e) {    std::cerr << "error loading the model\n";    return -1;  }  std::cout << "ok\n";}

两个文件编写完成之后便可对其执行编译：

mkdir example_testcd example_testcmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..cmake --example_test . --config Release

第四步：

给example-app.cpp添加模型推理代码并执行：

std::vector inputs;inputs.push_back(torch::ones({1, 3, 224, 224}));// 执行推理并将模型转化为Tensoroutput = module.forward(inputs).toTensor();std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

以上便是C++中部署PyTorch模型的全过程，相关教程可参考PyTorch官方：

https://pytorch.org/tutorials/

下载1：OpenCV-Contrib扩展模块中文版教程

在「小白学视觉」公众号后台回复：扩展模块中文教程，即可下载全网第一份OpenCV扩展模块教程中文版，涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。

下载2：Python视觉实战项目52讲

在「小白学视觉」公众号后台回复：Python视觉实战项目，即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目，助力快速学校计算机视觉。

下载3：OpenCV实战项目20讲

在「小白学视觉」公众号后台回复：OpenCV实战项目20讲，即可下载含有20个基于OpenCV实现20个实战项目，实现OpenCV学习进阶。

交流群

欢迎加入公众号读者群一起和同行交流，目前有SLAM、三维视觉、传感器、自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN、算法竞赛等微信群（以后会逐渐细分），请扫描下面微信号加群，备注：”昵称+学校/公司+研究方向“，例如：”张三 + 上海交大 + 视觉SLAM“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告，否则会请出群，谢谢理解~