polygraphy_tensorrt_engine_build_from_onnx_with_fp16_inference.py

1from polygraphy.backend.trt import CreateConfig, EngineFromNetwork, NetworkFromOnnxPath, TrtRunner
2from polygraphy.logger import G_LOGGER
3
4# 1. Define where the model is located
5onnx_path = "model.onnx"
6
7# 2. Build a TensorRT engine from the ONNX model
8# This uses the high-level Polygraphy loaders to simplify the process
9build_engine = EngineFromNetwork(NetworkFromOnnxPath(onnx_path), config=CreateConfig(fp16=True))
10
11# 3. Create a runner to manage the inference session
12# The context manager ensures that the engine and context are properly freed
13with TrtRunner(build_engine) as runner:
14    # 4. Run inference
15    # By default, the runner will generate random input data if none is provided
16    outputs = runner.infer()
17
18    # 5. Process the outputs
19    for name, tensor in outputs.items():
20        print(f"Output Name: {name} | Shape: {tensor.shape} | First 5 values: {tensor.flatten()[:5]}")