Turn PyTorch models into TensorRT engines by running them once

What it does

torch2trt takes a PyTorch module and sample input, runs a standard forward pass, and intercepts registered layer calls to build an equivalent TensorRT engine on the fly. The resulting TRTModule acts like the original model but executes optimized for NVIDIA GPUs, particularly Jetson Nano and Xavier devices. You can persist the compiled engine with a normal PyTorch state_dict.

The interesting bit

Rather than parsing a static graph or exporting to ONNX, the converter monkey-patches PyTorch functional calls—such as torch.nn.ReLU.forward—and attaches a TensorRT layer builder that piggybacks on live tensor data. You can extend coverage by writing a plain Python function and decorating it with @tensorrt_converter, which is unusually low-friction for a hardware compiler toolchain.

Key highlights

• Single-call conversion: torch2trt(model, [input]) returns an optimized module.
• Extensible in Python: register custom layer converters with a decorator without touching C++.
• Includes throughput benchmarks on Jetson Nano and Xavier for ResNet, VGG, DenseNet, and SqueezeNet.
• Optional plugin library handles layers that TensorRT does not support natively, compiled separately with CMake for better portability.
• Experimental community contributions include Quantization Aware Training support under torch2trt.contrib.

Caveats

• Coverage is explicitly limited; the tool was built primarily for the JetBot project and may fail on operators outside its tested set.
• Some unsupported layers require installing an additional C++ plugin library.
• Desktop GPU performance is not benchmarked in the README, so the documented gains are strictly edge-focused.

Verdict

Developers shipping PyTorch vision models on NVIDIA Jetson or edge GPUs should try this before reaching for ONNX or manual TensorRT APIs. If you are running transformers, exotic activations, or non-CNN architectures, expect to write custom converters or look elsewhere.