Skip to main content

LibTorch (PyTorch C++ Frontend)

LibTorch is the official C++ frontend for Pytorch. However, due to its lack of documentation, I encountered lots of confusions during its use. Some useful tips/tricks are listed here just FYI.

As of July 2021, this documentation is written based on my experience with libtorch v1.9

Performance: slower than Python​

It is repeatedly reported that inference using LibTorch is much slower than that in Python. See discussions in #19106.

There is also a ZhiHu article (in Chinese) that attempts to analyze this issue in-depth. The solution it proposed was to recompile libtorch by linking to libraries used by pytorch.

Cross-Save/Load Tensors in Python​

This section documents how to save tensors in C++ and load them into Python, and vice versa. It is often done for more friendly debugging experience offered by the Python frontend.

info

Note that LibTorch torch::save() function (source) saves the tensors in a wrapped TorchScript (JIT) module, unlike torch.save() (docs) in Python.

Save tensor in C++ and load in Python​

In C++, call torch::save() to save.

#include <torch/torch.h>

// save one tensor
torch::save(tensor, "tensor.pt");

// save multiple tensors
torch::save({tensora, tensorb, tensorc}, "tensors.pt");

In Python, use torch.jit.load() to load.

import torch

# Load one tensor
tensor_model = torch.jit.load("tensor.pt")
tensor = list(tensor_model.parameters())[0]

# Load multiple tensors
tensors_model = torch.jit.load("tensors.pt")
tensora = list(tensors_model.parameters())[0]
tensorb = list(tensors_model.parameters())[1]
tensorc = list(tensors_model.parameters())[2]

Save tensor in Python and load in C++​

The following codes are adapted from pytorch/pytorch#20356 (comment) and updated for the v1.8+ API (get_attribute => attr).

Save tensors in Python: to do so, you have to create a model and include all tensors into this TorchScript module.

import torch

class Container(torch.nn.Module):
    def __init__(self, my_values):
        super().__init__()
        for key in my_values:
            setattr(self, key, my_values[key])

my_values = {
    'a': torch.ones(2, 2),
    'b': torch.ones(2, 2) + 10,
    'c': 'hello',
    'd': 6
}

# Save arbitrary values supported by TorchScript
# https://pytorch.org/docs/master/jit.html#supported-type
container = torch.jit.script(Container(my_values))
container.save("container.pt")

Load tensors in C++

#include <torch/script.h>

torch::jit::script::Module container = torch::jit::load("container.pt");

torch::Tensor a = container.attr("a").toTensor();
torch::Tensor b = container.attr("b").toTensor();
std::string c = container.attr("c").toStringRef();

int64_t d = container.attr("d").toInt();

Alternative: use pickle​

An alternative is to use pickle_save and pickle_load (source). See this comment in pytorch/pytorch#20356 for usage.

Multiple Input/Output for Inference​

Suppose we have loaded a model named module and want to use it for inference. However, the model requires multiple inputs/outputs.

The codes are adapted from pytorch/pytorch#18337.

std::vector<torch::jit::IValue> inputs;
inputs.push_back(tensora);
inputs.push_back(tensorb);

auto outputs = module->forward(inputs).toTuple();

torch::Tensor out1 = outputs->elements()[0].toTensor();
torch::Tensor out2 = outputs->elements()[1].toTensor();
info

Note: if you only have one output, you can directly call toTensor() on the output of forward().