Porting Pytorch Models to C++ | Pipelines that Port Pytorch Models to C++

In the world of deep learning, Python frameworks like PyTorch have become the go-to tools for training neural networks due to their flexibility and ease of use. However, when it comes to deploying models in production applications, languages like C++ are often preferred for their speed and portability. Thankfully, it is possible to get the best of both worlds by training models in PyTorch and then porting them to C++ for efficient inference. In this article, we‘ll dive into the why and how of bringing your PyTorch models into the realm of C++.

Why Port PyTorch Models to C++?

There are several compelling reasons to run your machine learning models in C++ rather than Python:

1. Speed – C++ code generally executes much faster than Python. For applications that require low-latency, real-time inference, extracting the model from the Python interpreter overhead can yield significant speedups.

2. Portability – C++ is a highly portable language that can be compiled to run on a wide variety of operating systems and hardware architectures. This makes it easier to deploy your model as part of an application across different platforms.

3. Deployment – Python requires an interpreter to run the code whereas C++ is compiled ahead of time. Distributing your model as part of a compiled binary rather than a collection of Python scripts can make installation and deployment simpler for end users.

4. Integration – It‘s straightforward to integrate C++ code with other languages and systems. You can easily call your model from Java, Swift, or embedded environments like microcontrollers that may not support Python.

Overall, porting to C++ is a logical choice when you need to integrate a trained model into a larger application or system optimized for performance and scale. And with some handy tools from PyTorch, the process is not as daunting as it may seem.

The Porting Process

At a high level, the steps to port a PyTorch model to C++ are:

1. Train and save your model using PyTorch in Python.
2. Export the model to an intermediate format like TorchScript or ONNX.
3. Load the exported model in C++ and run inference.

The key to this process is using an interchange format that can capture the structure of the model in a way that can be executed by a C++ runtime. Two popular options for this are TorchScript and ONNX.

TorchScript

TorchScript is a representation of a PyTorch model that can be run in a high-performance C++ environment. It‘s a statically typed subset of Python that captures the structure of your model and provides a way to serialize it. There are two ways to create TorchScript:

1. Tracing – With tracing, you pass an example input through your model, and the sequence of operations performed is captured in a graph representation. This works well for models that have a static flow of control.

2. Scripting – For models that use more dynamic control flow like conditionals and loops, you can annotate your Python code with type hints to help the TorchScript compiler understand it. The model code is then explicitly parsed and compiled.

Here‘s an example of tracing a simple model and saving it:

import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(4, 5)

    def forward(self, x):
        return self.linear(x)

model = MyModel()
example_input = torch.rand(1, 4)

traced_script_module = torch.jit.trace(model, example_input)
traced_script_module.save("model.pt")

To run this in C++, you load the saved module and provide an input tensor:

#include <torch/script.h>

int main(int argc, char** argv) {
  torch::jit::script::Module module = torch::jit::load("model.pt");

  std::vector<torch::jit::IValue> inputs;
  inputs.emplace_back(torch::rand({1, 4}));

  at::Tensor output = module.forward(inputs).toTensor();
  std::cout << output << std::endl;
}

One nice feature of TorchScript is that it has a built-in optimizing compiler that can inline functions, eliminate dead code, and fuse operations to improve performance. Depending on the model, this can give a notable speed boost over running in Python.

ONNX

ONNX (Open Neural Network Exchange) is an open standard for representing machine learning models. Models in the ONNX format can be run using various engines and runtimes across different platforms. It‘s supported by many frameworks in addition to PyTorch including TensorFlow, scikit-learn, and XGBoost.

To convert a PyTorch model to ONNX, you use the torch.onnx.export function:

import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(4, 5)

    def forward(self, x):
        return self.linear(x)

model = MyModel()
example_input = torch.rand(1, 4)

torch.onnx.export(model, example_input, "model.onnx")

This saves the model in the ONNX protobuf format which specifies the computational graph, operator definitions, and model parameters. To run this in C++, you can use the ONNX Runtime library:

#include <onnxruntime_cxx_api.h>

int main(int argc, char** argv) {
  Ort::Env env;
  Ort::SessionOptions session_options;
  Ort::Session session(env, "model.onnx", session_options);

  std::vector<int64_t> input_shape = {1, 4};
  std::vector<float> input_data = {/* ... */};

  auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
  Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_data.data(), 
                             input_data.size(), input_shape.data(), input_shape.size());

  std::vector<Ort::Value> ort_outputs = session.Run(Ort::RunOptions{nullptr}, 
                                                    {"input"}, &input_tensor,
                                                    1, {"output"}, 1);

  float* output_data = ort_outputs[0].GetTensorMutableData<float>();
  std::cout << "Output: " << output_data[0] << std::endl;
}

The ONNX Runtime has some powerful features for optimizing model performance:

• It can leverage hardware acceleration libraries like Intel DNNL and NVIDIA TensorRT.
• It performs node-level optimizations like operator fusion.
• It can do quantization to further reduce model size and speed up execution.

So if you are deploying on GPU-equipped systems or resource-constrained devices, ONNX is definitely worth considering.

Performance Comparison

To see the kind of speedups you can expect when porting to C++, let‘s benchmark our example linear model running in PyTorch natively vs TorchScript vs ONNX Runtime. We‘ll run each for 1000 iterations on an Intel Core i5 CPU.

PyTorch:

$ python bench.py
PyTorch time: 0.16164779663085938 seconds

TorchScript:

$ python compile_script.py
$ ./bench_torchscript
TorchScript time: 0.09617900848388672 seconds

ONNX Runtime:

$ python convert_onnx.py 
$ ./bench_onnx 
ONNX time: 0.07968997955322266 seconds

We can see that TorchScript cuts the execution time almost in half compared to running in Python, and ONNX Runtime improves things a bit further. The relative performance will depend on the specific model architecture and the operations used. But in general, porting to C++ via TorchScript or ONNX will make your model noticeably snappier.

Optimization Tips

If you want to get the most performance out of your exported models in C++, here are a few things to keep in mind:

1. Batch your inputs – Running a model on a batch of data is usually more efficient than one example at a time due to better utilization of parallel hardware. If possible, structure your application to batch up inputs before invoking the model.

2. Fuse operations – Combining multiple operations into a single kernel can reduce memory overhead and unlock vectorization opportunities. Look for chances to use fused operations like Conv2d + ReLU or BatchNorm + ReLU.

3. Quantize weights – Neural networks can often be quantized to int8 without losing much accuracy. This reduces model size and takes advantage of high-throughput int8 instructions on modern CPUs. The ONNX Runtime has a quantization tool to help automate this process.

4. Use accelerator hardware – If deploying on a machine with a GPU or AI accelerator, make sure your runtime is configured to take advantage of it. The ONNX Runtime can plug into Intel‘s OpenVINO toolkit or NVIDIA‘s TensorRT to target specific inference hardware.

Real-World Examples

Porting PyTorch models to C++ has been a key part of deploying AI in many impactful real-world applications. Here are a few examples:

• Tesla uses PyTorch to train neural networks for autonomous driving tasks like object detection and motion prediction. These models are then exported via TorchScript and integrated into the self-driving software stack that runs on their custom inference hardware in the vehicle.

• Facebook uses ONNX to bridge between research in PyTorch and production services written in C++. Models for tasks like recommendation, ranking, and integrity are converted through the ONNX format to run on server hardware optimized for inference at scale.

• Researchers at MIT used TorchScript to deploy a deep learning model for real-time inverse kinematics on a Baxter robot arm. Porting to C++ reduced runtime latency enough to enable closed-loop control of the robot at 500Hz.

These examples highlight the importance of bridging between flexible research tools like PyTorch and high-performance production systems in C++ to bring the latest AI breakthroughs to the real world.

Conclusion

As we‘ve seen, porting PyTorch models to C++ is a powerful technique for integrating deep learning into production applications. By exporting through interchange formats like TorchScript or ONNX, you can take advantage of high-performance runtimes and hardware-specific optimizations to dramatically speed up model execution. With some tuning and careful deployment, PyTorch in C++ is a surefire way to scale up your AI systems to meet the demands of the real world. So go forth and deploy!