Increasing inference speed

In addition to running inference at reduced precision, Lightning Pose models can be accelerated further with additional techniques that speed up model processing – torch.compile(), ONNX Runtime, and TensorRT – and with hardware-accelerated video decoding – PyNvVideoCodec. This page benchmarks the model-processing techniques together (including eager FP16 for reference) and walks through how to use each one, assuming you’ve already trained a model with litpose train; see the PyNvVideoCodec section below for the video-decoding benchmarks and usage.

TL;DR

• All three techniques help, and complexity roughly tracks with payoff. torch.compile() is a one-line change and gives solid gains. ONNX Runtime requires an export step but is the simplest option to deploy without a full Python/PyTorch runtime. TensorRT requires the most setup (matching CUDA/TensorRT library versions) but wins across every model and GPU tested, up to 2.88x on the 6-view multi-view transformer on an L4 (TensorRT FP16 + DALI vs. eager FP32).

• None of these techniques change the model’s predictions. We compared final keypoint predictions against the eager FP32 baseline on real video frames for all three methods – max deviation was under 0.08px in every case, consistent with ordinary floating-point kernel differences and far below any dataset’s typical pixel error (see Accuracy check below).

Overview

• Eager FP16 / BF16 – no code changes beyond --precision, see Mixed Precision Training & Inference. Included in the table below for reference.

• torch.compile() – a single-line PyTorch feature that JIT-compiles the model’s forward pass into fused GPU kernels. No export step, no new dependencies. See Usage: torch.compile().

• ONNX Runtime – export the model to the ONNX format, then run inference through onnxruntime’s CUDAExecutionProvider. Lets you deploy without a full PyTorch install. See Usage: ONNX Runtime.

• TensorRT – same ONNX export, but run through onnxruntime’s TensorrtExecutionProvider, which builds an autotuned, hardware-specific inference engine. Most setup, biggest gains. See Usage: TensorRT.

• PyNvVideoCodec – hardware-accelerated video decoding via direct NVIDIA Video Codec SDK bindings, as an alternative to DALI’s video reader. Independent of the model-processing techniques above – see Usage: PyNvVideoCodec.

• These techniques are complementary and can be combined – e.g. torch.compile() or TensorRT for the model together with PyNvVideoCodec for decoding – for the best end-to-end throughput. See the figure in the following section.

Results

The figure below benchmarks every combination of model architecture, precision/runtime technique, video decoder, and GPU on real end-to-end litpose predict calls – not an isolated forward pass. Each bar is the mean of 10 timed repeats (1 discarded warmup run) for the two single-view panels; the 6-view multi-view panel uses 3 timed repeats (1 discarded warmup run) instead. Error bars show standard error of the mean.

End-to-end litpose predict timing for ResNet50 (single-view), ViT-S (single-view), and ViT-S (6-view multi-view). Each row shows 5 technique groups (eager FP32, eager FP16, torch.compile + FP16, ONNX Runtime FP16, TensorRT FP16), each with 4 bars: L4 + DALI, L4 + PyNvVideoCodec, A100 + DALI, A100 + PyNvVideoCodec.

TensorRT FP16 wins overall across every model and GPU, with the largest end-to-end gain (2.88x) on the 6-view multi-view transformer on L4 – notably larger than the same technique’s gain on A100, since A100’s much faster eager-FP32 baseline leaves less relative headroom to close. Decoder choice diverges sharply by model: for the single-view models, PyNvVideoCodec is often the faster decoder on A100 once any model-acceleration technique is applied (9-25% faster than DALI), though it trails DALI somewhat on L4. For the 6-view multi-view model, DALI wins outright across every technique and GPU – confirming the multi-view PyNvVideoCodec slowdown documented below carries through to full end-to-end timing, and the gap is substantially larger on A100 (19-36% slower) than on L4 (2-14% slower).

Caveats

• cuDNN TF32 was left on for the ResNet50 eager-FP32 baseline (matmul TF32 is off by PyTorch’s own default; predict_speed_matrix_benchmark.py sets neither flag explicitly, so both follow PyTorch’s defaults). This means the eager-FP32 numbers in the figure above are not a fully strict FP32 baseline either – doesn’t change the direction of any result here, since TensorRT/ONNX Runtime FP16 still win by a wide margin, but it does mean the eager-FP32 bars are a few percent faster than a true FP32 baseline would be for convolution-heavy models like ResNet50.

• A single dynamic-shape TensorRT profile was used per (model, GPU, precision), covering the full batch-size range tested, rather than a separate engine per exact batch size. This is simpler but can leave some performance on the table at batch sizes far from the profile’s optimum – for example, multi-view-FP32-on-L4 was roughly flat (0.91-1.07x) across batch sizes rather than showing a clear win.

• ONNX Runtime FP16 + DALI is missing from the figure for the 6-view multi-view model on L4, likely due to a race condition between DALI’s CUDA stream and the ONNX Runtime CUDAExecutionProvider on L4 GPUs – see issue 483.

Accuracy check

Before recommending any of these, we checked whether they change the model’s actual predictions. Using a real trained checkpoint (ResNet50, mirror-mouse-fused) and 5 sampled frames from a real test video, we compared keypoint predictions from each accelerated method against the eager FP32 reference, reusing the same real preprocessing and keypoint-extraction code as Model.predict_frame – the only thing swapped per method was the “images -> heatmaps” compute step itself.

Method	Mean pixel deviation vs. eager FP32	Max pixel deviation vs. eager FP32
Eager FP16	0.014px	0.049px
torch.compile (FP32)	0.004px	0.071px
ONNX Runtime (FP32)	0.003px	0.020px
TensorRT (FP32)	0.005px	0.078px

All four methods land within a small fraction of a pixel of the eager FP32 reference – consistent with expected floating-point kernel differences rather than any real change in the computation.

Usage

The tutorials below assume you’ve already trained a model with litpose train and have a model_dir containing config.yaml and a checkpoint.

torch.compile

Model.compile() compiles the model’s forward pass in place. Call it after loading the model and before running prediction:

from lightning_pose.api import Model

model = Model.from_dir("path/to/model_dir")
model.compile()

# use as normal -- the compiled graph is now used internally
result = model.predict_on_video_file("path/to/video.mp4")

The same thing is available from the CLI with the --compile flag:

litpose predict /path/to/model_dir /path/to/video.mp4 --compile

The first prediction after compiling triggers compilation and is therefore slower than an uncompiled run – expect tens of seconds of one-time overhead. This is expected. Subsequent calls with the same input shape reuse the compiled graph. Changing batch size or input resolution triggers a new compilation automatically. Calling compile() more than once is a no-op.

Note

compile() compiles the model’s forward method and assigns it back onto the model, rather than wrapping the whole module. Wrapping the module (model.model = torch.compile(model.model)) works if you call model.model(images) directly in your own code, but does not reliably engage the compiled graph when going through predict_on_video_file / predict_frame / predict_on_label_csv, because those call get_loss_inputs_labeled rather than forward – and calling anything other than forward/__call__ on a compiled module silently falls back to the original, uncompiled submodule.

Note

torch.compile() uses the TorchInductor backend, which generates Triton kernels and therefore requires a GPU of CUDA compute capability 7.0 or higher (Volta and newer). On older GPUs – for example the GTX 10-series (capability 6.x) – compile() raises a RuntimeError immediately, rather than failing partway through the first prediction. Check your device with python -c "import torch; print(torch.cuda.get_device_capability())".

ONNX Runtime

Installation

ONNX Runtime is an optional dependency. onnxruntime-gpu wheels are built against a specific CUDA major version at build time, and the wheel does not detect or adapt to whatever CUDA you already have – so check what your PyTorch install is using first:

python -c "import torch; print(torch.version.cuda)"

Then use that CUDA version to select the proper installation option from ONNX Runtime’s installation selector.

Exporting additionally requires the onnx and onnxscripts packages, which onnxruntime-gpu does not pull in as a dependency:

pip install onnx onnxscripts

Only Model.export() need them; a machine that just runs an already-exported .onnx file needs onnxruntime-gpu alone.

Usage

Exporting and loading are separate steps. Export once, then load with runtime="onnx" for every subsequent prediction run:

from lightning_pose.api import Model

# export once
model = Model.from_dir("path/to/model_dir")
model.export("onnx", onnx_precision="fp16")

# load through ONNX Runtime and predict as normal
onnx_model = Model.from_dir(
    "path/to/model_dir", runtime="onnx", onnx_precision="fp16"
)
result = onnx_model.predict_on_video_file("path/to/video.mp4")

The same two steps from the CLI:

litpose export /path/to/model_dir --runtime onnx --onnx-precision fp16
litpose predict /path/to/model_dir /path/to/video.mp4 --runtime onnx --onnx-precision fp16

Exports are written to a fixed location inside the model directory:

model_dir/
├── config.yaml
├── tb_logs/.../checkpoints/epoch=214-step=12685-best.ckpt
└── exports_onnx/
    ├── epoch=214-step=12685-best_fp16.onnx
    └── epoch=214-step=12685-best_fp32.onnx   # only if exported

You never pass export paths yourself – export() always writes here and from_dir(runtime="onnx") always reads from here. The filename is the checkpoint’s own stem plus the export precision, so it stays unambiguous which checkpoint an export came from if a model directory ever holds exports alongside an updated checkpoint. onnx_precision may be omitted when exactly one export exists for the checkpoint, and is required to disambiguate when several do.

Note

precision and onnx_precision are different settings.

• precision ("fp32"/"fp16"/"bf16", also litpose predict --precision) controls the autocast precision of an eager forward pass. It leaves the weights on disk untouched.

• onnx_precision ("fp32"/"fp16", also --onnx-precision) is the weight precision baked into the exported .onnx file itself.

With runtime="onnx", precision is ignored – the exported file’s own precision is what runs.

Note

runtime="onnx" rebinds the model’s forward method to the ONNX session rather than wrapping the module, for the same reason described in the torch.compile() note above.

Note

torch.compile() has no effect on an ONNX Runtime session. Combining --compile with --runtime onnx, or calling compile() on a model loaded with runtime="onnx", raises rather than silently doing nothing.

TensorRT

Installation

TensorRT and onnxruntime-gpu need matching major versions, and pip install tensorrt / pip install onnxruntime-gpu each default independently to whatever is newest – so they don’t necessarily agree with each other, or with your PyTorch install’s CUDA version. Two checks catch this before you ever call export("tensorrt", ...):

1. Find out what onnxruntime actually needs. Rather than trusting an external compatibility table that may lag your installed onnxruntime-gpu version, inspect the provider library it ships directly:

ldd "$(python -c "import onnxruntime, os; print(os.path.join(os.path.dirname(onnxruntime.__file__), 'capi', 'libonnxruntime_providers_tensorrt.so'))")" | grep libnvinfer

This prints the exact libnvinfer.so.<N> onnxruntime is linked against – <N> is the TensorRT major version you need (e.g. libnvinfer.so.10 means TensorRT 10.x; a newer TensorRT 11.x, which pip install tensorrt may give you by default, will not load).

2. Check what’s currently installed, and fix it if the major version doesn’t match:

python -c "import tensorrt; print(tensorrt.__version__)"

Install using the tensorrt-cu11/tensorrt-cu12/tensorrt-cu13 package matching your PyTorch CUDA version (python -c "import torch; print(torch.version.cuda)"), with an upper bound of (required major + 1) so pip picks the newest release of that major without jumping to the next one – e.g. if step 1 printed libnvinfer.so.10:

pip install "tensorrt-cu13<11"

For the full CUDA/cuDNN/TensorRT installation matrix, see NVIDIA’s TensorRT install guide.

3. Put the installed libraries on the loader’s search path. pip installs TensorRT’s .so files into site-packages/tensorrt_libs, which the system linker does not search by default. onnxruntime loads them with a raw dlopen, not through the tensorrt Python package – so importing tensorrt first does not help. TensorRT’s own libraries in turn dlopen their CUDA runtime dependencies (cuBLAS, cuDNN) the same raw way; if those were pip-installed as separate packages (as PyTorch’s CUDA wheels do), those directories need to be on the path too.

export LD_LIBRARY_PATH="$(python -c "import tensorrt_libs, os; print(os.path.dirname(tensorrt_libs.__file__))"):$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="$(dirname "$(find "$(python -c 'import site; print(site.getsitepackages()[0])')/nvidia" -name 'libcublas.so.*' | head -1)"):$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="$(dirname "$(find "$(python -c 'import site; print(site.getsitepackages()[0])')/nvidia" -name 'libcudnn.so.*' | head -1)"):$LD_LIBRARY_PATH"

cuBLAS and cuDNN are the two runtime libraries this needed beyond TensorRT’s own on a typical pip install; step 4 below will name any other missing library by its exact filename if your setup needs more – repeat the find-based pattern above, swapping in that filename, for each one it reports.

Add all of these to your shell profile (or the conda env’s activate.d script) so they persist across sessions rather than needing to be re-exported every time.

4. Verify the install. First, sanity-check that the tensorrt package itself loads its own bundled libraries:

python -c "import tensorrt as trt; assert trt.Builder(trt.Logger())"

This alone is not sufficient – it says nothing about whether onnxruntime’s independent dlopen of the same libraries will succeed, since onnxruntime never goes through the tensorrt Python package to load them. Test that separately:

python -c "import ctypes, onnxruntime, os; path = os.path.join(os.path.dirname(onnxruntime.__file__), 'capi', 'libonnxruntime_providers_tensorrt.so'); ctypes.CDLL(path); print('OK')"

If this raises OSError: libnvinfer.so.<N>: cannot open shared object file, either step 2 installed the wrong major version or LD_LIBRARY_PATH isn’t set in this shell.

Usage

TensorRT builds its engine from an existing ONNX export, so export to ONNX first if you haven’t already:

from lightning_pose.api import Model

model = Model.from_dir("path/to/model_dir")
model.export("onnx", onnx_precision="fp16")
model.export("tensorrt", onnx_precision="fp16", max_batch_size=32)

trt_model = Model.from_dir(
    "path/to/model_dir", runtime="tensorrt", onnx_precision="fp16"
)
result = trt_model.predict_on_video_file("path/to/video.mp4")

The same from the CLI:

litpose export /path/to/model_dir --runtime onnx --onnx-precision fp16
litpose export /path/to/model_dir --runtime tensorrt --onnx-precision fp16 --max-batch-size 32
litpose predict /path/to/model_dir /path/to/video.mp4 --runtime tensorrt --onnx-precision fp16

max_batch_size sets the upper end of the dynamic-shape batch profile the engine is built for – inference at a larger batch size fails. opt_batch_size (defaults to max_batch_size) is the batch size the engine is tuned to run fastest at; correctness holds for any batch size in [1, max_batch_size], but performance is best near opt_batch_size.

The engine cache is written alongside the ONNX exports:

model_dir/
├── config.yaml
├── tb_logs/.../checkpoints/epoch=214-step=12685-best.ckpt
├── exports_onnx/
│   └── epoch=214-step=12685-best_fp16.onnx
└── exports_trt/
    └── epoch=214-step=12685-best_fp16/
        ├── trt_metadata.json
        └── ...                          # engine cache files (managed by onnxruntime)

The first call to model.export("tensorrt", ...) builds and autotunes the engine, which can take anywhere from several seconds to a few minutes. trt_metadata.json records the GPU name, TensorRT/onnxruntime versions, and batch profile the engine was built with – Model.from_dir(..., runtime="tensorrt") reads this back and warns if the current GPU doesn’t match, since (unlike an .onnx file) a built engine is tied to the exact GPU architecture it was built on and is not portable across machines.

Note

Building always requires a real GPU with a working TensorRT/onnxruntime install – there is no CPU fallback tier for runtime="tensorrt". If TensorrtExecutionProvider can’t load, this raises rather than silently running on CUDAExecutionProvider or CPU (unlike plain ONNX Runtime, which does have a CPU tier – see the note in the ONNX Runtime section above about checking session.get_providers()).

Note

The accuracy check earlier on this page used TensorRT at FP32. At FP16, quantization can shift the heatmap’s predicted peak by a few pixels on keypoints the model has no real opinion about – a near-uniform heatmap, likelihood near the noise floor – a “peak-flipping” effect that happens under any precision or kernel change (it reproduces on the already-merged ONNX FP16 path too, given a real dataset rather than a handful of hand-picked frames) and isn’t specific to TensorRT. On confidently-tracked keypoints it stays small: repeated FP16 engine builds during testing measured mean deviation around 0.2-0.3px, though TensorRT’s own build-time autotuning adds some run-to-run variance on top – engine builds aren’t perfectly deterministic, so don’t expect bit-identical predictions from two separately-built engines either.

PyNvVideoCodec (Video Decoding)

Everything above speeds up the model’s forward pass. This section covers a separate, complementary axis: how the video frames themselves get decoded off disk before they ever reach the model. By default, litpose predict decodes video with DALI’s GPU video reader. PyNvVideoCodec (direct NVIDIA Video Codec SDK / NVDEC bindings) is an alternative decoder backend that talks to the hardware decoder more directly, skipping DALI’s pipeline/graph framework overhead.

Note

PreparePynvvc (the class backing --reader pynvvc) reads its window size from the existing cfg.dali config section – dali.base.predict.sequence_length / dali.context.predict.sequence_length – rather than a separate cfg.pynvvc section. The name is a little confusing, since this setting also governs how many frames PyNvVideoCodec reads per batch, but the windowing semantics (frames per iteration, overlap for context/MHCRNN models) are identical between the two reader backends, so one config value covers both rather than keeping two in sync. If you only ever use --reader pynvvc and never DALI, you still configure the window size via dali.*.predict.sequence_length.

Installation

pip install PyNvVideoCodec

Requires an NVIDIA GPU from the Turing generation or newer (Turing/Ampere/Ada/Hopper/Blackwell) and driver version 530.41.03 or newer on Linux. If PyNvVideoCodec can’t decode a given video on the current machine (unsupported GPU generation, driver too old, package not installed, or an unsupported video format), litpose predict automatically falls back to DALI (or, if DALI itself isn’t installed – e.g. on macOS/Windows, where nvidia-dali-cuda110 isn’t available – to OpenCV) rather than erroring – see below.

Usage

--reader is independent of --runtime/--compile above – it only controls how video frames are read, not how the model runs on them. From the CLI:

litpose predict /path/to/model_dir /path/to/video.mp4 --reader pynvvc
litpose predict /path/to/model_dir /path/to/video.mp4 --reader dali
litpose predict /path/to/model_dir /path/to/video.mp4 --reader opencv

Or from the API:

model = Model.from_dir("path/to/model_dir")
result = model.predict_on_video_file("path/to/video.mp4", reader="pynvvc")

Omitting --reader (or passing reader=None) auto-selects the best backend usable on the current machine for the given video: PyNvVideoCodec if it’s usable, else DALI if it’s installed, else OpenCV – no error, just a progressively slower/more-portable fallback. OpenCV decodes on CPU and needs no NVIDIA GPU or platform-specific package, so it’s the one backend guaranteed to work everywhere (including macOS, Windows, and GPU-less machines), making it the last rung of the fallback chain rather than a speed option. Explicitly requesting a specific --reader on a machine where it isn’t usable raises a clear error instead of silently falling back, so you know your run isn’t using the backend you asked for.