Click on the blue wordsFollow us

Due to changes in the public account’s push rules, please click “View” and add “Star” to get exciting technical shares as soon as possible

Source from the internet, infringement will be deleted

Introduction This article mainly explains how to deploy PyTorch models onto a C++ platform, detailing the model conversion, saving serialized models, loading serialized PyTorch models in C++, and executing Script Modules in four main sections.

Recently, due to work requirements, I need to deploy PyTorch models onto a C++ platform. The basic process mainly refers to the official teaching examples, during which I encountered many pitfalls, which I hereby record.

1. Model Conversion

Libtorch does not depend on Python; models trained in Python need to be converted to Script Modules to be loaded by libtorch for inference. In this step, the official website provides two methods: Method 1: Tracing This method is relatively simple. You only need to provide the model with a set of inputs, run through the inference network once, and then use torch.jit.trace to record the information along the path and save it. An example is as follows:

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)

The downside is that if the model contains control flow, such as if-else statements, one set of inputs can only traverse one branch, and in this case, it is impossible to record the complete model information. Method 2: Scripting Directly write the model in Torch script and annotate the model accordingly. Compile the module using torch.jit.script to convert it to ScriptModule. An example is 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) 

• The forward method will be compiled by default, and methods called within forward will also be compiled in the order they are called.
• If you want to compile a method that is not called by forward, you can add @torch.jit.export.
• If you want a method not to be compiled, you can use [@torch.jit.ignore](https://link.zhihu.com/?target=https%3A//pytorch.org/docs/master/generated/torch.jit.ignore.html%23torch.jit.ignore) or [@torch.jit.unused](https://link.zhihu.com/?target=https%3A//pytorch.org/docs/master/generated/torch.jit.unused.html%23torch.jit.unused).

# Same behavior as pre-PyTorch 1.2
@torch.jit.script
def some_fn():
    return 2

# Marks a function as ignored, if nothing
# ever calls it then this has no effect
@torch.jit.ignore
def some_fn2():
    return 2

# As with ignore, if nothing calls it then it has no effect.
# If it is called in script it is replaced with an exception.
@torch.jit.unused
def some_fn3():
  import pdb; pdb.set_trace()
  return 4

# Doesn't do anything, this function is already
# the main entry point
@torch.jit.export
def some_fn4():
    return 2

I encountered many pitfalls in this step, mainly due to the following two reasons:

1. Unsupported Operations

TorchScript supports a subset of Python operations. Most operations used in torch can find corresponding implementations, but there are some awkward unsupported operations. A detailed list can be found at unsupported-ops（https://pytorch.org/docs/master/jit_unsupported.html#jit-unsupported）， below are some operations I encountered:1) Variable number of parameters/return values are not supported, for example

def __init__(self, **kwargs):

or

if output_flag == 0:
    return reshape_logits
else:
    loss = self.loss(reshape_logits, term_mask, labels_id)
    return reshape_logits, loss

2) Various iteration operations eg1.

layers = [int(a) for a in layers] 

reports error torch.jit.frontend.UnsupportedNodeError: ListComp aren’t supported. It can be changed to:

for k in range(len(layers)):
    layers[k] = int(layers[k])

eg2.

seq_iter = enumerate(scores)
try:
    _, inivalues = seq_iter.__next__()
except:
    _, inivalues = seq_iter.next()

eg3.

line = next(infile)

3) Unsupported statements eg1. Unsupported continue torch.jit.frontend.UnsupportedNodeError: continue statements aren’t supported. eg2. Unsupported try-catch torch.jit.frontend.UnsupportedNodeError: try blocks aren’t supported. eg3. Unsupported with statement. 4) Other common ops/modules eg1. torch.autograd.Variable. Solution: use torch.ones/torch.randn etc. to initialize + .float()/.long() to specify data types. eg2. torch.Tensor/torch.LongTensor etc. Solution: same as above. eg3. requires_grad parameter is only supported in torch.tensor, not in torch.ones/torch.zeros etc. eg4. tensor.numpy() eg5. tensor.bool(). Solution: use tensor > 0 instead of tensor.bool(). eg6. self.seg_emb(seg_fea_ids).to(embeds.device). Solution: need to explicitly call .cuda() where the GPU is required. In summary: Try to avoid using libraries other than native Python and PyTorch, such as NumPy, and use various PyTorch APIs as much as possible.

2. Specifying Data Types

1) Attributes, most member data types can be inferred from their values, but empty lists/dictionaries need to be specified in advance:

from typing import Dict

class MyModule(torch.nn.Module):
    my_dict: Dict[str, int]

    def __init__(self):
        super(MyModule, self).__init__()
        # This type cannot be inferred and must be specified
        self.my_dict = {}

        # The attribute type here is inferred to be `int`
        self.my_int = 20

    def forward(self):
        pass

m = torch.jit.script(MyModule())

2) Constants, use _Final_ keyword:

try:
    from typing_extensions import Final
except:
    # If you don't have `typing_extensions` installed, you can use a
    # polyfill from `torch.jit`.
    from torch.jit import Final

class MyModule(torch.nn.Module):

    my_constant: Final[int]

    def __init__(self):
        super(MyModule, self).__init__()
        self.my_constant = 2

    def forward(self):
        pass

m = torch.jit.script(MyModule())

3) Variables. By default, they are tensor types and immutable, so non-tensor types must be specified:

def forward(self, batch_size:int, seq_len:int, use_cuda:bool):

Method three: Mixing Tracing and Scripting One is to call script in the trace model, suitable for cases where only a small part of the model needs to use control flow. The usage example is as follows:

import torch

@torch.jit.script
def foo(x, y):
    if x.max() > y.max():
        r = x
    else:
        r = y
    return r


def bar(x, y, z):
    return foo(x, y) + z

traced_bar = torch.jit.trace(bar, (torch.rand(3), torch.rand(3), torch.rand(3)))

Another case is to use tracing to generate submodules in the script module. For some layers that have Python features unsupported by the script module, the relevant layer can be encapsulated, and tracing can be used to record the flow of the relevant layer, while other layers do not need to be modified. An example usage is as follows:

import torch
import torchvision

class MyScriptModule(torch.nn.Module):
    def __init__(self):
        super(MyScriptModule, self).__init__()
        self.means = torch.nn.Parameter(torch.tensor([103.939, 116.779, 123.68])
                                        .resize_(1, 3, 1, 1))
        self.resnet = torch.jit.trace(torchvision.models.resnet18(),
                                      torch.rand(1, 3, 224, 224))

    def forward(self, input):
        return self.resnet(input - self.means)

my_script_module = torch.jit.script(MyScriptModule())

2. Saving the Serialized Model

If the pitfalls in the previous step are all avoided, saving the model is very simple. You just need to call save and pass a filename. It is important to note that if you want to train the model on GPU and do inference on CPU, you must convert it before saving the model, and remember to call model.eval(), as follows:

gpu_model.eval()
cpu_model = gpu_model.cpu()
sample_input_cpu = sample_input_gpu.cpu()
traced_cpu = torch.jit.trace(traced_cpu, sample_input_cpu)
torch.jit.save(traced_cpu, "cpu.pth")

traced_gpu = torch.jit.trace(traced_gpu, sample_input_gpu)
torch.jit.save(traced_gpu, "gpu.pth")

3. Loading the Trained Model in C++

To load serialized PyTorch models in C++, you must rely on the PyTorch C++ API (also known as LibTorch). Installing libtorch is very simple; you just need to download the corresponding version from the PyTorch official website and unzip it. You will get a folder structured as follows:

libtorch/
  bin/
  include/
  lib/
  share/

Then you can build the application. A simple example directory structure is as follows:

example-app/
  CMakeLists.txt
  example-app.cpp

The example code for example-app.cpp and CMakeLists.txt is as follows:

#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 ScriptModule from a file using 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";
}

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)

Now, you can run the following commands to build the application from the example-app/ folder:

mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
cmake --build . --config Release

Where /path/to/libtorch is the path where the previously downloaded libtorch folder is located. If this step is successful, you will see a 100% completion prompt. The next step is to run the compiled executable, and you will see the output “ok”. Congratulations!

4. Executing the Script Module

Finally, we are at the last step! Now you just need to pass inputs to the model and execute forward to get the output. A simple example is as follows:

// 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';

The first two lines create a vector of torch::jit::IValue and add a single input. Use torch::ones() to create the input tensor, equivalent to torch.ones in the C++ API. Then, run the script::Module‘s forward method and convert the returned IValue value to a tensor by calling toTensor(). C++ operations for torch are quite user-friendly, and you can find corresponding implementations by using torch:: or appending _ to the method, for example:

torch::tensor(input_list[j]).to(at::kLong).resize_({batch, 128}).clone()
//torch::tensor corresponds to pytorch's torch.tensor; at::kLong corresponds to torch.int64; resize_ corresponds to resize

Finally, check to ensure that the output on the C++ side is consistent with PyTorch, and you have successfully completed it! I have encountered countless pitfalls and lost countless hairs, and much of it was learned through trial and error. If there are any mistakes, please feel free to correct me!

5. References

[1] PyTorch C++ API – PyTorch master document: https://pytorch.org/cppdocs/[2] Torch Script – PyTorch master documentation: https://pytorch.org/tutorials/advanced/cpp_export.html

If you are over 18 years old and find learning 【C language】 too difficult? Want to try another programming language? I recommend you learn Python. The existing 499 yuan Python zero-basic course is limited to free access, with only 10 spots available!



▲ Scan the QR code - get it for free

Click to read the original text to learn more