Tripy: A Python Programming Model For TensorRT

Tripy is a Python programming model for TensorRT that aims to provide an excellent user experience without compromising performance. Some of the features of Tripy include:

  • Intuitive API: Tripy doesn’t reinvent the wheel: If you have used NumPy or PyTorch before, Tripy APIs should feel familiar.

  • Excellent Error Messages: When something goes wrong, Tripy tries to provide informative and actionable error messages. Even in cases where the error comes from deep within the software stack, Tripy is able to map it back to the Python code that caused it.

Installation

python3 -m pip install --no-index -f https://nvidia.github.io/TensorRT-Incubator/packages.html tripy --no-deps
python3 -m pip install -f https://nvidia.github.io/TensorRT-Incubator/packages.html tripy

Important: There is another package named tripy on PyPI. Note that it is not the package from this repository. Please use the instructions above to ensure you install the correct package.

Getting Started

We’ve included several guides in Tripy to make it easy to get started. We recommend starting with the Introduction To Tripy guide.

To get an idea of the look and feel of Tripy, let’s take a look at a short code example. All of the features used in this example are explained in more detail in the introduction guide mentioned above.

 1# Define our model:
 2class Model(tp.Module):
 3    def __init__(self):
 4        self.conv = tp.Conv(in_channels=1, out_channels=1, kernel_dims=[3, 3])
 5
 6    def __call__(self, x):
 7        x = self.conv(x)
 8        x = tp.relu(x)
 9        return x
10
11
12# Initialize the model and populate weights:
13model = Model()
14model.load_state_dict(
15    {
16        "conv.weight": tp.ones((1, 1, 3, 3)),
17        "conv.bias": tp.ones((1,)),
18    }
19)
20
21inp = tp.ones((1, 1, 4, 4))
22
23# Eager mode:
24eager_out = model(inp)
25
26# Compiled mode:
27compiled_model = tp.compile(
28    model,
29    args=[tp.InputInfo(shape=(1, 1, 4, 4), dtype=tp.float32)],
30)
31
32compiled_out = compiled_model(inp)

>>> model.state_dict()
{
    conv.bias: tensor([1.0000], dtype=float32, loc=gpu:0, shape=(1,)),
    conv.weight: tensor(
        [[[[1.0000, 1.0000, 1.0000],
           [1.0000, 1.0000, 1.0000],
           [1.0000, 1.0000, 1.0000]]]],
        dtype=float32, loc=gpu:0, shape=(1, 1, 3, 3)),
}
>>> inp
tensor(
    [[[[1.0000, 1.0000, 1.0000, 1.0000],
       [1.0000, 1.0000, 1.0000, 1.0000],
       [1.0000, 1.0000, 1.0000, 1.0000],
       [1.0000, 1.0000, 1.0000, 1.0000]]]],
    dtype=float32, loc=gpu:0, shape=(1, 1, 4, 4))
>>> eager_out
tensor(
    [[[[10.0000, 10.0000],
       [10.0000, 10.0000]]]],
    dtype=float32, loc=gpu:0, shape=(1, 1, 2, 2))
>>> compiled_model
Executable(x: Tensor) -> Tensor
>>> compiled_out
tensor(
    [[[[10.0000, 10.0000],
       [10.0000, 10.0000]]]],
    dtype=float32, loc=gpu:0, shape=(1, 1, 2, 2))

User Guides

API Reference

Developer Guides