README.md

Tripy: A Python Programming Model For TensorRT

Installation | Getting Started | Documentation | Notebooks | Examples | Contributing

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

• 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 usually able to map it back to the related Python code.

• Friendly Documentation: The documentation is meant to be accessible and comprehensive, with plenty of examples to illustrate important points.

Installation

Installing Prebuilt Wheels

python3 -m pip install nvtripy -f https://nvidia.github.io/TensorRT-Incubator/packages.html

Building Wheels From Source

To get the latest changes in the repository, you can build Tripy wheels from source.

1. Make sure build is installed:

   python3 -m pip install build

2. From the tripy root directory, run:

   python3 -m build . -w

3. Install the wheel, which should have been created in the dist/ directory. From the tripy root directory, run:

   python3 -m pip install -f https://nvidia.github.io/TensorRT-Incubator/packages.html dist/nvtripy-*.whl

4. [Optional] To ensure that Tripy was installed correctly, you can run a sanity check:

   python3 -c "import nvtripy as tp; x = tp.ones((5,), dtype=tp.int32); assert x.tolist() == [1] * 5"

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.

Other features covered in our guides include:

• Compiling code (including dynamic shape support)
• Quantization

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.

# Define our model:
class Model(tp.Module):
    def __init__(self):
        self.conv = tp.Conv(in_channels=1, out_channels=1, kernel_dims=[3, 3])

    def __call__(self, x):
        x = self.conv(x)
        x = tp.relu(x)
        return x


# Initialize the model and populate weights:
model = Model()
model.load_state_dict(
    {
        "conv.weight": tp.ones((1, 1, 3, 3)),
        "conv.bias": tp.ones((1,)),
    }
)

inp = tp.ones((1, 1, 4, 4))

# Eager mode:
eager_out = model(inp)

# Compiled mode:
compiled_model = tp.compile(
    model,
    args=[tp.InputInfo(shape=(1, 1, 4, 4), dtype=tp.float32)],
)

compiled_out = compiled_model(inp)