This tutorial has been updated for use with PyTorch 1.2

As the name implies, PyTorch’s primary interface is the Python programming language. While Python is the preferred language for many scenarios that require dynamics and ease of iteration, there are also many situations in which these properties work against Python. The latter usually applies in an environment that requires production – low latency and tight deployment. For production scenarios, C ++ is often the language of choice, even if it is only bound to another language like Java, Rust, or Go. The following paragraphs outline the path PyTorch provides from an existing Python model to a serialized representation that can be loaded and executed entirely from C ++, without relying on Python.

Step 1: Convert the PyTorch model to a Torch script

The PyTorch model’s journey from Python to C ++ is initiated by Torch Script, a representation of the PyTorch model that can be understood, compiled, and serialized by the Torch Script compiler. If you started with an existing PyTorch model written using the Vanilla “Eager” API, you must first convert the model to a Torch script. In the most common cases (described below), this takes very little effort. If you already have the Torch script module, you can skip to the next part of this tutorial.

There are two ways to convert the PyTorch model into a Torch script. The first is called tracing, a mechanism that captures the structure of a model by evaluating it once using sample inputs and recording the flow of those inputs through the model. This applies to models with limited use of control flow. The second method is to add explicit comments to the model to tell the Torch Script compiler that the model code can be parsed and compiled directly based on the constraints imposed by the Torch Script language.

Tip: You can find full documentation on both methods, along with further guidance on how to use them, in the official Torch Script Reference.

Method 1: Convert to a Torch script by tracing

To convert the PyTorch model into a Torch script by tracing, you must pass an instance of the model along with the sample input to the function Torch.jit.trace. This will produce a torch.jit.ScriptModule object whose model evaluation trace will be embedded in the module’s forward method:

Import the torch import torchvision # you an instance of the model. The model = torchvision. Models. Resnet18 (#) you are usually provided to model forward () method of the sample input. example = torch.rand(1, 3, 224, Traced_script_module = torch. Jit. Trace (model, example)Copy the code

Now you can evaluate the traced ScriptModule to be the same 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

Method 2: Convert to a Torch script with comments

In some cases, for example, if the model has a particular form of control flow, you may need to write the model directly in the Torch script and annotate the model accordingly. For example, suppose you have the following Vanilla 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 outputCopy the code

Because the forward method of this module uses an input dependent control flow, it is not suitable for tracing. Instead, we can convert it to ScriptModule. To convert the module to ScriptModule, you need to compile the module using torch.jit. Script as follows:

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

If you need to exclude methods from nn.Module because they use Python features not yet supported by TorchScript, you can annotate them with @torch.jit. Ignore

My_module is an instance of ScriptModule and can be serialized.

Step 2: Serialize the script module to a file

Once you have a ScriptModule (by tracing or annotating the PyTorch model), you can serialize it as a file. Later, you will be able to use C ++ to load the module from this file and execute it without relying on Python. Suppose we want to serialize the ResNet18 model shown earlier in the trace example. To do this serialization, simply call save on the module and pass a filename:

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

This will generate the traced_resnet_model.pt file in your working directory. If you also want to serialize my_module, call my_module.save(“my_module_model.pt”). We have now officially left the Python realm and are ready to step into the C ++ realm.

Step 3: Load the script module in C ++

To load the serialized PyTorch model in C ++, your application must rely on the PyTorch C ++ API (also known as LibTorch). The LibTorch distribution contains a collection 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. For this tutorial, we will use CMake and LibTorch to build a minimal C ++ application that simply loads and executes the serialized PyTorch model.

Smallest C ++ application

Let’s start by discussing the code for loading modules. The following will have been done:

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 script module from the file with the following command: 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"; }Copy the code

The header contains all the relevant contains in the LibTorch library needed to run the example. Our application takes the serialized PyTorch ScriptModule file path as its only command-line argument, and then proceeds to deserialize the module using the Torch :: Jit :: Load () function, which takes the file path as input. As return, we received a Torch: : jit: : script: : Module object. We will discuss how to implement it later.

Depends on LibTorch and building the application

Suppose we store the above code in a file named example-app.cpp. The smallest cmakelists.txt might look simple:

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 final thing to do to set up the sample application is the LibTorch distribution. You can always get the latest stable version from the Download page on the PyTorch website. If you download and unzip the latest archive, you should receive a folder with the following directory structure:

libtorch/
  bin/
  include/
  lib/
  share/Copy the code

• The lib/ folder contains shared libraries that you must link to,
• Include/folder contains header files that the program needs to include,
• The share/ folder contains the necessary CMake configuration to enable the simplicity abovefind_package(Torch)Command.

Tip; On Windows, debug and distribution versions are not ABI compatible. If you plan to build your project in debug mode, try using the debug version of LibTorch.

The final step is to build the application. To do this, assume the following layout for the sample directory:

example-app/
  CMakeLists.txt
  example-app.cppCopy the code

Now we can run the following command to build the application from example-app/ folder:

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

/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 provide the trace ResNet18 model traced_resnet_model.pt path we created earlier to the sample application binaries, we should be rewarded with a friendly “OK.” Note that if you try to run this example using my_module_model.pt, you will receive an error message indicating that the shapes you entered are incompatible. My_module_model.pt requires 1D instead of 4D.

root@4b5a67132e81:/example-app/build# ./example-app <path_to_model>/traced_resnet_model.pt
okCopy the code

Step 4: Execute the script module in C ++

Having successfully loaded the serialized ResNet18 with C ++, we now only need to execute a few lines of code! Let’s add these lines to the C ++ application’s main() function:

STD ::vector<torch::jit::IValue> inputs; inputs.push_back(torch::ones({1, 3, 224, 224})); At ::Tensor output = module. Forward (inputs).totensor (); // At ::Tensor output = module. std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';Copy the code

The first two lines set the inputs for our model. We created a torch:: JIT ::IValue vector (erased in the Script::Module method of type type-erased) and added a single input. To create the input tensor, we use torch::ones(), equivalent to torch.ones in the C ++ API. We then run the forward method of Script ::Module and pass it the input vector we created. In return, we get a new IValue, which we transform into a tensor by calling toTensor().

Tip: to overall understanding about the torch: : ‘ones and PyTorch C + + API functions for more information, please refer to the documentation, https://pytorch.org/cppdocs. The PyTorch C ++ API provides almost the same functional parity as the Python API, allowing you to further manipulate and manipulate tensors as you would in Python.

In the last line, we print out the first five entries. Since we provided the same input to the 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 traced_resnet_model.pt-0.2698-0.0381 0.4023-0.3010-0.0448 [ Variable [CPUFloatType] {1, 5}]Copy the code

For reference, Python’s previous output was:

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

Looks like a good match!

Tip: To move the model to GPU memory, write model.to (at::kCUDA); . Ensure that the input to the model is also in CUDA memory by calling tensor.to (at::kCUDA), which will return the new tensor in CUDA memory.

Step 5: Get help and explore the API

Hopefully, this tutorial will give you an overview of the PyTorch model’s path from Python to C ++. Using the concepts described in this tutorial, you should be able to go from vanilla, the “Eager” PyTorch model, to compiled ScriptModules in Python, to serialized files on disk, and — end of loop — to modules in executable script: C ++.

Of course, there are many concepts that we haven’t covered. For example, you might find yourself wanting to extend a ScriptModule with a custom operator implemented in C ++ or CUDA and execute that custom operator in a ScriptModule loaded into a pure C ++ production environment. The good news is: it’s possible, and it’s well supported! You can now browse through the examples in this folder, and we will provide a tutorial shortly. For now, the following links may usually be helpful:

• The Torch Script reference: https://pytorch.org/docs/master/jit.html
• PyTorch C + + API documentation: https://pytorch.org/cppdocs/
• PyTorch Python API documentation: https://pytorch.org/docs/

As always, if you have any questions or concerns, please use our forum or GitHub Issues to get in touch.

Rock and the AI technology blog resources summary station: http://docs.panchuang.net/PyTorch, the official Chinese tutorial station: Chinese official document: http://pytorch.panchuang.net/OpenCV http://woshicver.com/