tensorrt-optimization

Inference optimization for deep models. The flow takes a trained torch network, exports it to ONNX, and applies graph and precision optimizations. The TensorRT backend is optional. When the TensorRT execution provider is missing the code falls back to the ONNX Runtime CPU provider, which is the path the tests exercise so everything runs offline on any machine.

Why this exists

When you ship a model the training framework graph is rarely the fastest way to run it. A common production recipe is to freeze the model to ONNX, let a runtime fold and fuse the graph, drop the weights to a lower precision, and hand the result to an accelerated backend like TensorRT. This repo is a small, honest implementation of that recipe with the parts that need real hardware kept optional.

What the pipeline does

1. Export. A torch module is traced in eval mode and written to ONNX. The batch dimension is marked dynamic so the exported model accepts any batch size at inference. The written graph is validated with the ONNX checker.

2. Graph optimization. ONNX Runtime runs its full optimization level over the graph. This folds constants, removes redundant nodes, and applies fusions such as conv plus batchnorm folding and conv plus activation fusion. The optimized graph is serialized so it can be reloaded and inspected.

3. Precision optimization. The graph is optionally converted to float16. Every float32 weight is cast to float16 and the model inputs and outputs are bridged with cast nodes so callers still pass and receive float32. When the onnxconverter-common package is installed it does this conversion. When it is not, a compact built in converter in src/optimize.py does the same job, so the precision pass works with no extra dependency.

4. Backend selection. The execution provider list is built in preference order. TensorRT comes first when present, then CUDA, then CPU. ONNX Runtime assigns each node to the first provider that can run it, so the same code uses TensorRT on a machine that has it and the CPU provider everywhere else.

5. Verification. The optimized ONNX model is run on the same input as the torch model and the largest absolute difference is measured. Graph optimization should leave the numerics unchanged within float noise. Float16 trades a little precision for speed, so it is checked against a looser tolerance.

Layout

src/
  models.py     tiny torch models used as optimization targets
  export.py     torch to ONNX export
  optimize.py   graph optimization, float16 conversion, sessions, runners
  backend.py    execution provider selection and TensorRT detection
  verify.py     numerical comparison between torch and ONNX
tests/
  test_export.py     export validity and dynamic batch behavior
  test_optimize.py   graph and precision passes and the report
  test_verify.py     optimized model matches torch within tolerance
  test_backend.py    provider ordering and the optional TensorRT path

The models are deliberately small. They stand in for the large pretrained networks you would optimize in practice, but the architecture is real: a convolutional net with a residual block and batchnorm gives the graph optimizer genuine fusions to perform, and the tiny size keeps the tests fast and offline.

Usage

import torch
from src.models import TinyConvNet
from src.export import export_to_onnx
from src.optimize import optimize_onnx_model, build_session, run_onnx

model = TinyConvNet().eval()
sample = torch.randn(1, 3, 16, 16)

export_to_onnx(model, sample, "model.onnx")

# Graph optimization only.
report = optimize_onnx_model("model.onnx", "model.opt.onnx", precision="fp32")
print(report.as_dict())

# Graph optimization plus float16.
report = optimize_onnx_model(
    "model.onnx", "model.opt.onnx", precision="fp16", fp16_path="model.fp16.onnx"
)

# Run the optimized model. TensorRT is used if available, otherwise CPU.
session = build_session("model.fp16.onnx", prefer_tensorrt=True)
out = run_onnx(session, sample.numpy())

Running the tests

python -m pytest tests/ -q

The suite uses tiny synthetic tensors, no downloads, and no API keys. It covers export validity, dynamic batch support, that graph optimization does not increase node count or change outputs, that the float16 pass produces float16 weights and a valid graph, that provider selection always falls back to CPU, and that the optimized model matches the torch model within tolerance.

The TensorRT path

TensorRT and the GPU build of ONNX Runtime are listed as optional in requirements.txt and are not installed here. src/backend.py detects the TensorRT execution provider at runtime. On a machine that has it, set prefer_tensorrt=True and the session uses it with no code change. The float16 artifact this pipeline produces is exactly the input a TensorRT backend wants, so the optimization work done on CPU carries straight over.