Translator: talengu

PyTorch’s primary interface is Python. While Python has the advantages of dynamic programming and ease of iteration, in many cases it is these attributes of Python that work against it. We often encounter production environments that meet low latency and stringent deployment requirements. For production scenarios, C++ is often the preferred language and can easily be bundled to another language, such as Java, Rust, or Go. This tutorial will walk you through the process from serializing the model representation of PyTorch training to loading _ and executing _ in C++.

Step 1: Convert the PyTorch model to Torch Script

The conversion of the PyTorch model from Python to C++ is implemented by Torch Script. Torch Script is a representation of the PyTorch model that is understood, compiled, and serialized by the Torch Script compiler. If you use the PyTorch model written with the base “Eager” API, you must first convert the model to Torch Script, which is also relatively easy. If you already have the Torch Script for the model, you can skip to the next part of the tutorial.

There are two ways to convert the PyTorch model to Torch Script. The first method is Tracing. This approach captures the structure of the model by evaluating the process by feeding samples into the model once. And record the flow of the sample in the model. This method is suitable for models in which control flow is rarely used. The second method is to add explicit annotations to the model, notifying the Torch Script compiler that it can parse and compile model code directly, subject to the constraints imposed by the Torch Script language.

Tips can be found in the official Torch Script reference for full documentation of both methods, as well as detailed guidance on which method to use.

Convert the model to Torch Script using Tracing

To convert the PyTorch model into a Torch script using tracing, an instance of the model and sample input must be passed to the Torch.jit. Trace function. This will generate a torch.jit.ScriptModule object and embed a trace of the model evaluation in the module’s forward method:

Import torch import torchvision # access model instance model. = torchvision models. Resnet18 # () to generate a sample for the network to spread before forward example = () torch.rand(1, 3, 224, Traced_script_module = torch. Jit. Trace (model, example)Copy the code

Now, the traced ScriptModule can perform the same calculations as the regular PyTorch module:

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>)Copy the code

Convert the Model to Torch Script via annotations

In some cases, for example, if the model uses a particular form of control flow, you might want to write the model directly in Torch Script and annotate the model accordingly. For example, suppose you have the following generic Pytorch model:

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
Copy the code

Because the forward method of this module uses inputting dependent control flow, it is not suitable to generate Torch Script using the Tracing method. To do this, we can convert a model to ScriptModule by inheriting torch.jit.ScriptModule and adding the @torch.jit. Script_method annotation to the model’s forward 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 input.sum() > 0:
          output = self.weight.mv(input)
        else:
          output = self.weight + input
        return output

my_script_module = MyModule()
Copy the code

Creating a new MyModule object now directly generates a serializable ScriptModule instance.

Step 2: Serialize the Script Module to a file

Either way, you can serialize a ScriptModule to a file, and C++ can execute the Pytorch model to which the Script corresponds without relying on any Python code. Suppose we want to serialize the ResNet18 model shown in the previous Trace example. To do this serialization, simply call save on the module and give it a filename:

traced_script_module.save("model.pt")
Copy the code

This will generate a model.pt file in the working directory. You can now leave Python behind and be ready to make the leap to C ++ language calls.

Step 3: load your Script Module in C++

To load the serialized PyTorch model in C ++, the application must rely on the PyTorch C ++ API – also known as _LibTorch_. The _LibTorch distribution _ contains a set of shared libraries, header files, and CMake build profiles. Although CMake does not rely on LibTorch’s requirements, it is the recommended method and will be well supported in the future. In this tutorial, we will use CMake and LibTorch to build a minimal C++ application that loads and executes the serialized PyTorch model.

Smallest C++ application

The following can be done to load modules:

#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"; }Copy the code

The

header contains all the relevant includes in the LibTorch library needed to run the example. The file path as its main function is to accept serialization ScriptModule only command line parameters, and then use the torch: : jit: : the load () function deserialization Module, get a pointer to the torch: : jit: : script: : the Module of Shared pointer, Equivalent to the torch. Jit.ScriptModule object in C ++. Finally, we only verify that this pointer is not null. We show how to implement it next.

Rely on the LibTorch library and build the application

Let’s save the above code in a file named example-app.cpp. The corresponding simple cmakelists.txt to build it is:

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)Copy the code

The last thing we did to build the sample application was download the LibTorch distribution. Get the latest stable version download Page from the Download page of the PyTorch website. If you download and unzip the latest archive, you have the following directory structure:

libtorch/
  bin/
  include/
  lib/
  share/
Copy the code
  • lib/Contains shared libraries with links,
  • include/Inclusion program needsincludeHeader file,
  • share/Include the necessary CMake configuration files to makefind_package(Torch) 。

Tips On Windows, debug and release builds are not ABI-compatible. To use debug, compile the PyTorch method using source code.

The final step is to build the application. To do this, assume our sample directory layout looks like this:

example-app/
  CMakeLists.txt
  example-app.cpp
Copy the code

We can now run the following command to build the application from example-app/ folder:

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

Where /path/to/libtorch should be the full path of the unzipped libtorch distribution. If all goes well, it will look something like this:

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-appCopy the code

If we provided the previous path to serialize the ResNet18 model to example-app, C++ would output OK:

root@4b5a67132e81:/example-app/build# ./example-app model.pt
ok
Copy the code

Run Script Module in C++ code

After successfully loading our serialized ResNet18 in C++, we add a few more lines of execution code to our C++ application’s main() function:

// 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';
Copy the code

The first two lines set up the inputs to our model. Creates a torch:: JIT ::IValue (a data type accepted and returned by script::Module objects) vector and adds an input. To create the input tensor, we use torch::ones() (C++ API) as in python torch.ones. Then we run script::Module’s forward method, pass in the input vector we created, and return a new IValue, which we’ll translate into tensors by calling toTensor().

Tips more on the corresponding C++ apis for torch::ones and PyTorch pytorch.org/cppdocs. The PyTorch C++ API is similar to the Python API, allowing you to handle tensors as Python does.

In the last line, we print out the first five entries. Since we provided the same input to our model in Python earlier in this tutorial, ideally we should see the same output. Let’s try it by recompiling our application and running it with the same serialization model:

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}]Copy the code

For reference, the output from the previous Python code is:

Tensor ([-0.2698, -0.0381, 0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)Copy the code

C++ output is the same as Python output. Success!

Tip To put your model on the GPU, you can write model->to(at::kCUDA); . To (at::kCUDA). This will return a new tensor in CUDA.

Step 5: Advanced tutorial and detailed API

This tutorial is intended to help you understand how the PyTorch model is called from python to c++. With the tutorial above, you can make a simple model by “eager” PyTorch, convert it to a ScriptModule, and serialize it. Then load the runtime model from script::Module in C++.

Of course, there’s a lot we haven’t covered. For example, if you want to implement custom actions in ScriptModule in C++ or CUDA, you can then run the ScriptModule model in C++ calls. This can be done, refer to this. Here are some more helpful documents:

  • Torch Script Reference: pytorch.org/docs/master…
  • PyTorch C++ API documentation: pytorch.org/cppdocs/
  • PyTorch Python API documentation: pytorch.org/docs/

If you have any bugs or problems, contact Pytorch Forum or Pytorch GitHub Issues for help.