Torch-TensorRT

In-framework compilation of PyTorch inference code for NVIDIA GPUs

Torch-TensorRT is a inference compiler for PyTorch, targeting NVIDIA GPUs via NVIDIA’s TensorRT Deep Learning Optimizer and Runtime. It supports both just-in-time (JIT) compilation workflows via the torch.compile interface as well as ahead-of-time (AOT) workflows. Torch-TensorRT integrates seamlessly into the PyTorch ecosystem supporting hybrid execution of optimized TensorRT code with standard PyTorch code.

More Information / System Architecture:

• Torch-TensorRT 2.0

Getting Started

• Installation

User Guide

• Torch-TensorRT Explained

• Dynamic shapes with Torch-TensorRT

• Post Training Quantization (PTQ)

• Saving models compiled with Torch-TensorRT

• Deploying Torch-TensorRT Programs

• DLA

• Compile Mixed Precision models with Torch-TensorRT

Tutorials

• Torch Compile Advanced Usage

• Deploy Quantized Models using Torch-TensorRT

• Engine Caching

• Engine Caching (BERT)

• Refitting Torch-TensorRT Programs with New Weights

• Serving a Torch-TensorRT model with Triton

• Torch Export with Cudagraphs

• Overloading Torch-TensorRT Converters with Custom Converters

• Using Custom Kernels within TensorRT Engines with Torch-TensorRT

• Mutable Torch TensorRT Module

• Weight Streaming

• Pre-allocated output buffer

Dynamo Frontend

• TensorRT Backend for torch.compile

• Compiling Exported Programs with Torch-TensorRT

TorchScript Frontend

• Creating a TorchScript Module

• Using Torch-TensorRT in Python

• Using Torch-TensorRT in C++

• Using Torch-TensorRT TorchScript Frontend Directly From PyTorch

FX Frontend

• Torch-TensorRT (FX Frontend) User Guide

Model Zoo

• Compiling ResNet with dynamic shapes using the torch.compile backend

• Compiling BERT using the torch.compile backend

• Compiling Stable Diffusion model using the torch.compile backend

• Compiling GPT2 using the Torch-TensorRT torch.compile frontend

• Compiling GPT2 using the dynamo backend

• Compiling Llama2 using the dynamo backend

• Compiling SAM2 using the dynamo backend

• Legacy notebooks

Python API Documentation

• torch_tensorrt

• torch_tensorrt.dynamo

• torch_tensorrt.logging

• torch_tensorrt.fx

• torch_tensorrt.ts

• torch_tensorrt.ts.ptq

C++ API Documentation

• Namespace torch_tensorrt

• Namespace torch_tensorrt::logging

• Namespace torch_tensorrt::ptq

• Namespace torch_tensorrt::torchscript

CLI Documentation

• torchtrtc

Contributor Documentation

• System Overview

• Writing Dynamo Converters

• Writing Dynamo ATen Lowering Passes

• Writing TorchScript Converters

• Useful Links for Torch-TensorRT Development

Indices

• Operators Supported

• Index

• Search Page

Legacy Further Information (TorchScript)

• Introductory Blog Post

• GTC 2020 Talk

• GTC 2020 Fall Talk

• GTC 2021 Talk

• GTC 2021 Fall Talk

• PyTorch Ecosystem Day 2021

• PyTorch Developer Conference 2021

• PyTorch Developer Conference 2022