下面几段将概述PyTorch提供的从现有Python模型到序列化表示的路径,该序列化表示完全可以从c++加载和执行,不依赖于Python。
Torch Script 是Pytorch模型的一种表示,可以被Torch Script编译器所理解,编译和序列化。在多数情况下,将模型转化为Torch Script只需要很少的工作。 将PyTorch模型转换为Torch脚本有两种方法。第一种方法称为追踪(tracing),这是一种机制,通过使用示例输入对模型进行一次评估,并记录这些输入在模型中的流动,从而捕获模型的结构。这适用于有限使用控制流的模型。第二种方法是向模型添加显式注释,通知Torch脚本编译器它可以直接解析和编译您的模型代码,这取决于Torch脚本语言施加的约束。
要通过跟踪将PyTorch模型转换为Torch脚本,您必须将模型的一个实例和一个示例输入传递给 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)追踪的ScriptModule现在可以与常规的PyTorch模块进行相同的计算:
In[1]: output = traced_script_module(torch.ones(1, 3, 224, 224)) In[2]: output[0, :5] Out[2]: tensor([-0.2698, -0.0381, 0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)在某些情况下,例如,如果您的模型使用特定形式的控制流,您可能希望直接用Torch脚本编写模型,并相应地对模型进行注释。例如,假设您有以下朴素的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由于该模块的正向方法使用依赖于输入的控制流,不适合追踪。相反,我们可以通过从torch.jit子类化ScriptModule来将其转换为ScriptModule。ScriptModule并添加@torch.jit.script_method注释到模型的正向方法:
import torch class MyModule(torch.jit.ScriptModule): def __init__(self, N, M): super(MyModule, self).__init__() self.weight = torch.nn.Parameter(torch.rand(N, M)) @torch.jit.script_method def forward(self, input): if bool(input.sum() > 0): output = self.weight.mv(input) else: output = self.weight + input return output my_script_module = MyModule(2, 3)现在,创建一个新的MyModule对象将直接生成ScriptModule的一个实例,该实例已准备好进行序列化。
现在,创建一个新的MyModule对象将直接生成ScriptModule的一个实例,该实例已准备好进行序列化。一旦掌握了ScriptModule(通过跟踪或注释PyTorch模型),就可以将其序列化到文件中。稍后,您将能够用c++从该文件加载模块,并在不依赖于Python的情况下执行它。假设我们想序列化前面跟踪示例中显示的ResNet18模型。要执行此序列化,只需调用模块上的save并传递一个文件名:
traced_script_module.save("model.pt")这将在工作目录中生成model.pt文件。我们现在已经正式离开了Python的领域,并准备跨越到c++的领域。
要用c++加载序列化的PyTorch模型,应用程序必须依赖于PyTorch c++ API——也称为LibTorch。LibTorch发行版包含一系列共享库、头文件和CMake构建配置文件。虽然CMake不是依赖于LibTorch的必要条件,但它是推荐的方法,将来会得到很好的支持。在本教程中,我们将使用CMake和LibTorch构建一个最小的c++应用程序,它只加载并执行一个序列化的PyTorch模型。
让我们从讨论加载模块的代码开始。下面就可以了:
#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; } // Deserialize the ScriptModule from a file using torch::jit::load(). std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(argv[1]); assert(module != nullptr); std::cout << "ok\n"; }假设我们将上面的代码存储到一个名为example-app.cpp的文件中。构建它的最小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 11)构建示例应用程序所需的最后一件事是LibTorch发行版。您可以从PyTorch网站的下载页面获取最新的稳定版本。如果你下载并解压最新的档案文件,你应该会收到一个目录结构如下的文件夹:
libtorch/ bin/ include/ lib/ share/ lib/文件夹包含必须链接的共享库,include/文件夹包含程序需要包含的头文件,share/文件夹包含必要的CMake配置,以启用上面的简单find_package(Torch)命令。最后一步是构建应用程序。为此,假设我们的示例目录如下所示:
example-app/ CMakeLists.txt example-app.cpp我们现在可以运行以下命令来构建应用程序内的例子-app/文件夹:
mkdir build cd build cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch .. make其中/path/to/libtorch应该是解压缩的libtorch发行版的完整路径。如果一切顺利,它将是这样的:
root@4b5a67132e81:/example-app# mkdir build root@4b5a67132e81:/example-app# cd build root@4b5a67132e81:/example-app/build# cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch .. -- The C compiler identification is GNU 5.4.0 -- The CXX compiler identification is GNU 5.4.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 -- Configuring done -- Generating done -- Build files have been written to: /example-app/build root@4b5a67132e81:/example-app/build# make Scanning dependencies of target example-app [ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o [100%] Linking CXX executable example-app [100%] Built target example-app如果我们将之前创建的序列化ResNet18模型的路径提供给最终的示例-app二进制文件,应该会得到一个友好的“ok”:
root@4b5a67132e81:/example-app/build# ./example-app model.pt ok在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(一个类型擦除的值类型script::Module方法接受和返回)的向量,并添加一个输入。为了创建输入张量,我们使用torch::ones(),它等价于torch.ones。然后运行脚本::Module的forward方法,将创建的输入向量传递给它。作为返回,我们得到一个新的IValue,我们通过调用toTensor()将其转换为一个张量。
在最后一行中,我们打印输出的前五个条目。由于我们在本教程的前面用Python为模型提供了相同的输入,所以理想情况下,我们应该看到相同的输出。让我们重新编译我们的应用程序,并运行它与相同的序列化模型:
root@4b5a67132e81:/example-app/build# make Scanning dependencies of target example-app [ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o [100%] Linking CXX executable example-app [100%] Built target example-app root@4b5a67132e81:/example-app/build# ./example-app model.pt -0.2698 -0.0381 0.4023 -0.3010 -0.0448 [ Variable[CPUFloatType]{1,5} ]作为参考,之前Python的输出为:
tensor([-0.2698, -0.0381, 0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)看起来输出相符!
小贴士 要将模型移动到GPU内存,可以写model->to(at::kCUDA);。通过调用tensoro .to(at::kCUDA),确保存在于CUDA内存中的模型的输入也在CUDA内存中,它将返回CUDA内存中的一个新的张量。
本教程希望让您对PyTorch模型从Python到c++的路径有一个大致的了解。有了本教程中描述的概念,您应该能够从普通的“eager”PyTorch模型,到Python中编译的ScriptModule,到磁盘上的序列化文件,再到c++中的可执行script::Module。
当然,有许多概念我们没有涉及。例如,您可能想要使用一个用c++或CUDA实现的自定义操作符来扩展ScriptModule,并在纯c++生产环境中加载的ScriptModule中执行这个自定义操作符。好消息是:这是可能的,并且得到了很好的支持!现在,您可以查看这个文件夹中的示例,稍后我们将提供一个教程。现时,以下连结可能会有帮助:
The Torch Script reference: https://pytorch.org/docs/master/jit.htmlThe PyTorch C++ API documentation: https://pytorch.org/cppdocs/The PyTorch Python API documentation: https://pytorch.org/docs/一如既往,如果您遇到任何问题或有任何疑问,您可以使用我们的论坛或GitHub问题来取得联系。