Click the above “Beginner’s Guide to Vision” and select “Star” or “Pin”
Important content delivered promptly
This article mainly explains how to deploy a PyTorch model on a C++ platform, detailing the model conversion, saving serialized models, loading serialized PyTorch models in C++, and executing Script Modules in four main sections.>> Join the Extreme City CV technology group to stay at the forefront of computer vision.
Recently, due to work requirements, I needed to deploy a PyTorch model on a C++ platform. The basic process mainly referred to the official teaching examples, during which I discovered many pitfalls, which I am recording here.
1. Model Conversion
Libtorch does not depend on Python; models trained in Python need to be converted to a Script Model to be loaded and inferred by libtorch. The official website provides two methods for this step:
Method 1: Tracing
This method is relatively simple, requiring just an input to the model to traverse the inference network, then using 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 there are control flows in the model, such as if-else statements, one input can only traverse one branch, making it impossible to fully record the model information in such cases.
Method 2: Scripting
Directly write the model in Torch script and annotate it accordingly, then compile the module using torch.jit.script to convert it into a 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 in forward, you can add @torch.jit.export.
If you want a method not to be compiled, you can use @torch.jit.ignore (https://pytorch.org/docs/master/generated/torch.jit.ignore.html#torch.jit.ignore) or @torch.jit.unused (https://pytorch.org/docs/master/generated/torch.jit.unused.html#torch.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
In this step, I encountered many pitfalls, mainly due to the following two reasons:
1. Unsupported Operations
The operations supported by TorchScript are a subset of Python; most operations used in torch can find corresponding implementations, but there are also some awkward unsupported operations. A detailed list can be found at https://pytorch.org/docs/master/jit_unsupported.html#jit-unsupported. Here are some operations I encountered:
1) Parameters/Return values do not support variable numbers, such as:
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]
Results in torch.jit.frontend.UnsupportedNodeError: ListComps 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. continue is not supported.
torch.jit.frontend.UnsupportedNodeError: continue statements aren’t supported.
eg2. try-catch is not supported.
torch.jit.frontend.UnsupportedNodeError: try blocks aren’t supported.
eg3. with statements are not supported.
4) Other common ops/modules
eg1. torch.autograd.Variable
Solution: Use torch.ones/torch.randn, etc., initialization + .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; torch.ones/torch.zeros are not available.
eg4. tensor.numpy()
eg5. tensor.bool()
Solution: Replace tensor.bool() with tensor > 0.
eg6. self.seg_emb(seg_fea_ids).to(embeds.device)
Solution: Explicitly call .cuda() where GPU conversion is needed.
In summary: try to avoid using libraries other than native Python and PyTorch, like NumPy, and use various PyTorch APIs as much as possible.
2. Specifying Data Types
1) Attributes: Most member data types can be inferred from 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 the 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 of tensor type and immutable, so non-tensor types must be specified.
def forward(self, batch_size:int, seq_len:int, use_cuda:bool):
Method 3: Mixing Tracing and Scripting
One way is to call script within a trace model, suitable for cases where only a small part of the model needs control flow. An 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)))
The other way is to use tracing to generate submodules within a script module for layers that have Python features unsupported by the script module, encapsulating the relevant layers and using trace to record the flow of those layers without modifying other layers. An example 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 Serialized Models
If you have navigated through the pitfalls in the previous step, saving the model is very simple; just call save and pass a filename. Note that if you want to train the model on GPU and do inference on CPU, you must convert it before saving, 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 Trained Models 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; just download the corresponding version from the PyTorch official website (https://pytorch.org/) and unzip it. You will get a folder structured as follows.
libtorch/
bin/
include/
lib/
share/
Then you can build your 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)
At this point, you can run the following commands to build the application from the <span>example-app/</span> folder:
mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
cmake --build . --config Release
Where /path/to/libtorch is the path to the downloaded libtorch folder. If this step goes smoothly, you will see a 100% completion message. The next step is to run the generated executable, and you should see the output “ok”. Congratulations!
4. Executing Script Module
Finally, we reach the last step! Just pass the constructed input 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 <span>torch::jit::IValue</span> inputs and add a single input. Using <span>torch::ones()</span> creates an input tensor, equivalent to <span>torch.ones</span> in the C++ API. Then, run the <span>script::Module</span>’s <span>forward</span> method, converting the returned IValue to a tensor using <span>toTensor()</span>. C++ operations in torch are quite friendly; you can find corresponding implementations by using <span>torch::</span> or adding _ at the end, such as:
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 matches that of PyTorch, and you will have successfully completed the task!
I encountered countless pitfalls and lost many hairs, and much of this was discovered through trial and error. If there are any errors, please feel free to correct me!
References
PyTorch C++ API – PyTorch master document
Torch Script – PyTorch master documentation
Good news!
The Beginner's Guide to Vision knowledge group is now open to the public👇👇👇
Download 1: OpenCV-Contrib Extension Module Chinese Version Tutorial
Reply: Extension Module Chinese Tutorial in the background of “Beginner's Guide to Vision” official account to download the first OpenCV extension module tutorial in Chinese on the internet, covering installation of extension modules, SFM algorithms, stereo vision, target tracking, biological vision, super-resolution processing, etc., with more than twenty chapters of content.
Download 2: Python Vision Practical Project 52 Lectures
Reply: Python Vision Practical Project in the background of “Beginner's Guide to Vision” official account to download 31 visual practical projects, including image segmentation, mask detection, lane line detection, vehicle counting, eyeliner addition, license plate recognition, character recognition, emotion detection, text content extraction, facial recognition, etc., to help quickly learn computer vision.
Download 3: OpenCV Practical Project 20 Lectures
Reply: OpenCV Practical Project 20 Lectures in the background of “Beginner's Guide to Vision” official account to download 20 practical projects based on OpenCV, achieving advanced learning of OpenCV.
Group Chat
Welcome to join the reader group of the official account to communicate with peers. Currently, there are WeChat groups for SLAM, 3D vision, sensors, autonomous driving, computational photography, detection, segmentation, recognition, medical imaging, GAN, algorithm competitions, etc. (these will be gradually subdivided). Please scan the WeChat number below to join the group, and note: “Nickname + School/Company + Research Direction”, for example: “Zhang San + Shanghai Jiao Tong University + Visual SLAM”. Please follow the format; otherwise, entry will not be allowed. After successful addition, invitations to related WeChat groups will be sent based on research direction. Please do not send advertisements in the group; otherwise, you will be removed. Thank you for your understanding~