Skip to content

README.md

Latest commit

 

History

History
154 lines (120 loc) · 7.22 KB

README.md

File metadata and controls

154 lines (120 loc) · 7.22 KB
Raw

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.

Measured results

bench/ contains a real benchmark on real hardware: a ResNet50 fine tuned on the NLM Montgomery chest X-ray set, run at five configurations on an RTX 5070 Ti.

configuration b8 latency b8 throughput speedup vs CPU accuracy TB recall agreement w/ FP32
CPU FP32 39.16 ms 204 img/s 1.00x 0.8571 0.7222 1.0000
CUDA FP32 4.42 ms 1812 img/s 8.87x 0.8333 0.7222 0.9762
TensorRT FP32 2.69 ms 2969 img/s 14.53x 0.8333 0.7222 0.9762
TensorRT FP16 1.13 ms 7091 img/s 34.71x 0.8333 0.7222 0.9762
TensorRT INT8 1.21 ms 6601 img/s 32.31x 0.8333 0.6667 0.9286

FP16 gives 34.7x over CPU for no changed predictions. INT8 is slower than FP16 below batch 32 and costs minority class recall, which the identical aggregate accuracy column hides entirely.

Full writeup, including the four silent failure modes that made the TensorRT provider report itself available while running every op on CPU: bench/README.md.

C++ pipeline, CUDA kernels and Nsight profile

cpp/ runs the same model and test images end to end in C++: TensorRT engines built with the TensorRT C++ API, a custom CUDA preprocessing path, and NVTX ranges traced with Nsight Systems.

TensorRT FP16, batch 32 resize and normalise end to end per image throughput
CPU, 1 thread 7.418 ms 7.756 ms 129 img/s
CPU, 8 threads 1.342 ms 1.474 ms 678 img/s
CUDA kernels 0.107 ms 0.579 ms 1726 img/s
  • The CUDA kernels reproduce Pillow's antialiased resize exactly: 0 of 2,107,392 pixels differ from the torchvision preprocessing the model was trained on.
  • Inference was never the bottleneck. At 0.1 ms per image it was 75 times cheaper than resizing on one core. With the kernels in place the next bottleneck is copying a 20 MB radiograph over PCIe, 69% of the GPU pipeline, and the Nsight trace shows exactly that.
  • The C++ FP32 engine with TF32 disabled agrees with PyTorch CPU FP32 on all 42 predictions, confirming the TF32 explanation in bench/README.md.

Details, build steps and the profile: cpp/README.md.

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

bench/          measured precision benchmark on chest X-rays (Python, ONNX Runtime TensorRT provider)
cpp/            C++ TensorRT pipeline, CUDA preprocessing kernels, tests, Nsight traces
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

src/backend.py detects the TensorRT execution provider at runtime, so the library code above runs anywhere and uses TensorRT when it is present. The measured runs use a real TensorRT install: ONNX Runtime's TensorRT provider in bench/, and the TensorRT C++ API directly in cpp/. The tests in tests/ exercise the CPU fallback so they run offline on any machine.