How to Deploy a Model with C++: Exporting Models With PyTorch Lightning

This is Part 1 of How to Deploy a Model with C++.

I am a big fan of C++ and Machine Learning and I enjoy playing around with Libtorch, but if you are going to start building a project, I wouldn’t suggest that you go about doing the ML part entirely in C++.  The main reason is that developing models in Python is a much faster process and there are many tools available such as PyTorch Lightning and Hydra that make working with Python much more convenient.  Not to mention, the parts of machine learning that really require optimization are typically already written in C++ and just have a Python wrapper around them.

However, there is one thing you may want to use C++ for, which is inference. In that situation, rather than building and training a model in C++, what you actually want to do is take a model created using PyTorch in Python and then deploy it in a C++ environment.  This allows you to get the best of both worlds.

Let’s look at three ways that we can export our models from Python:

1. Export a PyTorch model as a Checkpoint.
   
2. Export a PyTorch model in Onnx format.
   
3. Export a PyTorch model in TorchScript format

Each method has different advantages and disadvantages and we will discuss those in more depth as we cover each one.  I will also cover how to load these models in C++ in Part 2 of this article.

To start off, I’ll reuse the PyTorch Lightning model based on a Resnet18 from the simple PyTorch Lightning article I wrote previously.  First train the model using the FasionMNIST data module and take note of the final performance.  We can use that in Part 2 of this article when we import our exported models to make sure they are working the way we expect them to.

Export the Weights of a PyTorch Model

The first and most straightforward way of exporting a model is a checkpoint.  PyTorch Lightning creates a checkpoint automatically for you at the end of training.  You can also use the ModelCheckpoint callback if you want more control over when checkpoints are made.  If you haven’t specified where logs are saved, the PyTorch Lightning will save checkpoints in the `lightning_logs` directory.

Checkpoints are essentially the equivalent of an entire model saved in Pytorch. The main advantage of this format is that you don’t need anything special to load it afterwards.  Everything can be done with PyTorch Lightning and then on the C++ in Libtorch.

Export a Model in Onnx Format

The next method for exporting a model is by exporting it to the Onnx format.  Onnx is an accelerator that can be run on lots of different platforms.  In addition to being a convenient way to export a model from PyTorch Lightning, in some cases Onnx can be faster at inference than PyTorch and it comes with extensive documentation on improving performance.

Conveniently, PyTorch Lightning comes with a simple “to_onnx” function and that is all it takes to export your model. One thing to note is that you do need a sample input, but you can easily whip one up with

# (batch_size, channels, height, width)
input_sample = torch.ones(1, 3, 28, 28)
model.to_onnx("export.onnx", input_sample, export_params=True)

Export a Model in TorchScript Format

TorchScript is PyTorch’s own way of exporting a model in a more optimized way to be run in other environments such as C++.  It is just as simple to create as an Onnx model.  In this case all you have to do is use the “to_torchscript” function and there is no need for a sample input, so it is even easier than the “to_onnx” function.

model.to_torchscript(file_path="export.pt")

Additionally, if you pass the parameter “method=’trace’”, you can export a model in trace format instead of script format, but this is something you should research before trying it.

Conclusion

Developing in Python and then exporting your models to be used in a different environment is a great way to get the benefits of developing in Python while also taking advantage of being able to deploy in more specialized environments like C++ or on mobile phones.  Now that we have three different ways of exporting models from PyTorch Lightning, we can look at how we would actually import them in a C++ environment.  We will talk about that in Part 2 of this section.