【官方教程】C++平台用libtorch加载pytorch训练好的模型并推理
目录
1.1 方法1:通过tracing(跟踪)转换为Torch Script
1.2 方法2:通过annotation(注释)转换为Torch Script
3.3 最小的C ++应用程序:example-app.cpp
本文主要来自官方手册(https://pytorch.org/tutorials/advanced/cpp_export.html,pytorch 1.5.1版本),并添加了我的补充内容。本文输出结果为我实测结果,可能与官方文档不同。
顾名思义,PyTorch的主要接口是Python编程语言。尽管Python是合适于许多需要动态性和易于迭代的场景,并且是首选的语言,但同样的,在许多情况下,Python的这些属性恰恰是不利的。后者通常适用的一种环境是production(生产) ——要求低延迟和严格部署的地方。对于生产场景,即使只将C ++绑定到Java、Rust或Go之类的另一种语言中,C ++也是经常选择的语言。以下各段将概述PyTorch提供的从现有Python模型到可以完全从C ++加载和执行的序列化表示形式的路径,而无需依赖Python。
一、将PyTorch模型转换为Torch Script
PyTorch模型从Python到C ++的旅程由Torch Script启动,Torch Script是PyTorch模型的一种表示形式,可以由Torch Script编译器理解,编译和序列化。如果您是从使用vanilla“eager” API编写的现有PyTorch模型开始的,则必须首先将模型转换为Torch Script。在最常见的情况下(如下所述),这只需要花费很少的功夫。如果您已经有了Torch Script模块,则可以跳到本教程的下一部分。
有两种将PyTorch模型转换为Torch Script的方法。第一种称为tracing(跟踪),它是一种机制,通过使用示例输入对模型的结构进行一次评估,并记录这些输入在模型中的流量,从而捕获模型的结构。这适用于有限使用控制流的模型。第二种方法是在模型中添加显式批注,以告知Torch Script编译器可以根据Torch Script语言施加的约束直接解析和编译模型代码。
提示:您可以在官方Torch Script参考中找到有关这两种方法的完整文档,以及使用方法的进一步指导。
1.1 方法1:通过tracing(跟踪)转换为Torch Script
要将PyTorch模型通过跟踪转换为Torch Script,必须将模型的实例以及示例输入传递给torch.jit.trace函数。这将产生一个torch.jit.ScriptModule对象,该对象的模型评估的trace将嵌入模块的forward方法中:
import torch
import torchvision
# 你模型的一个实例.
model = torchvision.models.resnet18()
# 您通常会提供给模型的forward()方法的示例输入。
example = torch.rand(1, 3, 224, 224)
# 使用`torch.jit.trace `来通过跟踪生成`torch.jit.ScriptModule`
traced_script_module = torch.jit.trace(model, example)
现在可以对跟踪的ScriptModule进行评估,使其与常规PyTorch模块相同:
output = traced_script_module(torch.ones(1, 3, 224, 224))
output[0, :5]
Out[28]: tensor([-0.0877, -0.5285, -0.2473, 0.6111, 0.4302], grad_fn=<SliceBackward>)
1.2 方法2:通过annotation(注释)转换为Torch Script
在某些情况下,例如,如果模型采用特定形式的控制流,则可能需要直接在Torch Script中编写模型并相应地注释模型。例如,假设您具有以下vanilla Pytorch模型:
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
因为此模块的前向方法使用取决于输入的控制流,所以它不适合跟踪。相反,我们可以将其转换为ScriptModule。为了将模块转换为ScriptModule,需要使用torch.jit.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
my_module = MyModule(10,20)
sm = torch.jit.script(my_module)
如果您需要在nn.Module中排除某些方法,因为它们使用了TorchScript尚不支持的Python功能,则可以使用@torch.jit.ignore对其进行注释。
my_module是ScriptModule的实例,可以序列化。
二、将脚本模块序列化为文件
一旦有了ScriptModule(通过跟踪或注释PyTorch模型),您就可以将其序列化为文件了。稍后,您将可以使用C ++从此文件加载模块并执行它,而无需依赖Python。假设我们要序列化先前在跟踪示例中显示的ResNet18模型。要执行此序列化,只需在模块上调用save并传递一个文件名即可:
import torch
import torchvision
# 你模型的一个实例.
model = torchvision.models.resnet18()
# 您通常会提供给模型的forward()方法的示例输入。
example = torch.rand(1, 3, 224, 224)
# 使用`torch.jit.trace `来通过跟踪生成`torch.jit.ScriptModule`
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("traced_resnet_model.pt")
这将在您的工作目录中生成traced_resnet_model.pt文件。如果您还想序列化my_module,请调用my_module.save("my_module_model.pt")
我们现在已经正式离开Python领域,并准备跨入C ++领域。
三、在C ++中加载脚本模块
要在C ++中加载序列化的PyTorch模型,您的应用程序必须依赖于PyTorch C ++ API(也称为LibTorch)。LibTorch发行版包含共享库,头文件和CMake构建配置文件的集合。虽然CMake不是依赖LibTorch的一个必需品,但它是推荐的方法,并且将来会得到很好的支持。 对于本教程,我们将使用CMake和LibTorch构建一个最小的C ++应用程序,该应用程序简单地加载并执行序列化的PyTorch模型。
3.1 下载libtorch
建立示例应用程序的最后一件事是需要LibTorch发行版。您可以随时从PyTorch网站的下载页面上获取最新的稳定版本,如下图所示。
pytorch和libtorch的版本一定要保持一致,因为我pytorch的版本为1.5.1,所以我libtorch也下载的1.5.1版本。(如果需要其它的版本,则复制下方的链接,修改后面的数字为相应的版本号即可。)
可以先下载debug版本的,这样方便调试。
提示;在Windows上,debug 和release 版本不兼容。 如果您打算以debug 模式构建项目,请尝试使用LibTorch的debug 版本。

下载下来之后解压缩到下图中的文件夹中。

如果下载并解压缩最新的归档文件,则应收到具有以下目录结构的文件夹:
libtorch/
bin/
include/
lib/
share/
-
lib/ 文件夹包含您必须链接的共享库,
-
include/ 文件夹包含程序需要包含的头文件,
-
share/ 文件夹包含必要的CMake配置,以启用下面的简单
find_package(Torch)命令。
然后我们给系统添加环境变量:
C:\thirdLib\libtorch\lib
3.2 构建应用程序
假定示例目录的布局如下:
cmake_bin/
example-app/
CMakeLists.txt
example-app.cpp
traced_resnet_model.pt
3.3 最小的C ++应用程序:example-app.cpp
让我们从讨论加载模块的代码开始。以下将已经做:
#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 {
// 使用以下命令从文件中反序列化脚本模块: 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";
}
<torch/script.h>标头包含运行示例所需的LibTorch库中的所有相关包含。我们的应用程序接受序列化的PyTorch ScriptModule的文件路径作为其唯一的命令行参数,然后使用torch::jit::load()函数继续对该模块进行反序列化,该函数将此文件路径作为输入。作为返回,我们收到一个Torch::jit::script::Module对象。我们将稍后讨论如何执行它。
3.4 CMakeLists.txt
假设我们将以上代码存储在名为example-app.cpp的文件中。最小的CMakeLists.txt可能看起来很简单:
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)
set(Torch_DIR C:/thirdLib/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)
3.5 CMake构建项目
如果你使用的是libtorch(cpu)版,则你在generating之后,界面应类似下图。

如果你使用的是libtorch(gpu)版,则你在generating之后,界面应类似下图。
注意:如果要使用libtorch 1.5.1 gpu版,须安装cudatoolkit和cudnn。anaconda的Python环境给了我们很好地搭配pytorch 1.5.1 的cudatoolkit和cudnn的版本的示例,所以在C++环境中,我们的cudatoolkit和cudnn的版本可以和Python环境保持一致,分别为cudatoolkit10.2.89+cudnn7.6.5。
下图中提示“CUDA_SDK_ROOT_DIR”没有找到,这个问题暂未解决。不过实测对后续代码的运行没有影响。

四、在C ++中执行脚本模块
4.1 模型及模型的输入在CPU内存
成功用C ++加载了序列化的ResNet18之后,我们现在只需执行几行代码即可!让我们把以下几行代码添加到上面C ++应用程序的main()函数中的最后:
// 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';
前两行设置了我们模型的输入。我们创建一个torch::jit::IValue的向量(类型为type-erased的值Script::Module方法接受并返回),并添加单个输入。要创建输入张量,我们使用torch::ones(),等效于C ++ API中的torch.ones。然后,我们运行script::Module的forward方法,并向其传递我们创建的输入向量。作为回报,我们得到一个新的IValue,通过调用toTensor()将其转换为张量。
提示:要总体上了解有关torch::ones和PyTorch C ++ API之类的功能的更多信息,请参阅其文档,网址为https://pytorch.org/cppdocs。
PyTorch C ++ API提供了与Python API几乎相同的功能奇偶校验,使您可以像在Python中一样进一步操纵和处理张量。
在最后一行中,我们打印输出的前五个条目。由于在本教程前面的部分中,我们向Python中的模型提供了相同的输入,因此理想情况下,我们应该看到相同的输出。让我们通过重新编译我们的应用程序并以相同的序列化模型运行它来进行尝试。
注意:命令行参数为保存的.pt格式的模型文件的绝对路径,注意这个路径要包含文件名。
运行结果如下:

请注意,如果尝试使用my_module_model.pt运行此示例,则会收到一条错误消息,提示您输入的形状不兼容。my_module_model.pt需要1D而不是4D。
作为参考,Python以前的输出为:

看来匹配得很好!
4.2 将模型及模型的输入移至GPU内存
提示:要将模型移至GPU内存,可以编写model.to(at::kCUDA);。通过调用tensor.to(at::kCUDA),确保模型的输入也位于CUDA内存中,这将在CUDA内存中返回新的张量。
添加了将模型移至GPU内存、将模型的输入移至GPU内存的语句之后,最终的程序如下:
#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 {
// 使用以下命令从文件中反序列化脚本模块: torch::jit::load().
module = torch::jit::load(argv[1]);
}
catch (const c10::Error& e) {
std::cerr << "error loading the model\n";
return -1;
}
module.to(at::kCUDA); //将模型移至GPU内存
std::cout << "ok\n";
// Create a vector of inputs.
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({ 1, 3, 224, 224 }).to(at::kCUDA));//确保模型的输入在GPU内存
// 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';
}
输出如下:

可以看到输出的结果也在GPU内存中。
五、获取帮助并探索API
本教程有望使您对PyTorch模型从Python到C ++的路径有一个大致的了解。使用本教程中描述的概念,您应该能够从vanilla, “eager” PyTorch模型,
到Python中的已编译ScriptModule,再到磁盘上的序列化文件,以及–结束循环–到可执行脚本: C ++中的模块。
当然,有许多我们没有介绍的概念。例如,您可能会发现自己想要使用以C ++或CUDA实现的自定义运算符扩展ScriptModule,并在加载到纯
C ++生产环境中的ScriptModule中执行此自定义运算符。好消息是:这是可能的,并且得到了很好的支持!现在,您可以浏览此文件夹中的示例,
我们将很快提供一个教程。 目前,以下链接通常可能会有所帮助:
-
Torch Script参考:https://pytorch.org/docs/master/jit.html
-
PyTorch C ++ API文档:https://pytorch.org/cppdocs/
-
PyTorch Python API文档:https://pytorch.org/docs/
与往常一样,如果您遇到任何问题或疑问,可以使用我们的论坛或GitHub issues进行联系。
魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。
更多推荐


所有评论(0)