aitune.torch.backend.tensorrt.tensorrt_builder

TensorRT Builder module for building TensorRT engines from ONNX models.

Module Contents

Classes

Name	Description
`TensorRTBuilder`	Class for building TensorRT engines from ONNX models.

Data

logger

trt

API

class aitune.torch.backend.tensorrt.tensorrt_builder.TensorRTBuilder(
    input_onnx_path: pathlib.Path,
    output_path: pathlib.Path,
    workspace_size: int | None = None,
    optimization_level: int | None = None,
    compatibility_level: int | None = None,
    timing_cache: pathlib.Path | None = None,
    profiles: list[polygraphy.backend.trt.Profile] | None = None,
    enable_tf32: bool = True
)

Class for building TensorRT engines from ONNX models.

profiles

= profiles or []

aitune.torch.backend.tensorrt.tensorrt_builder.TensorRTBuilder._build_create_config_kwargs(
    max_workspace_size,
    optimization_level,
    compatibility_level,
    profiles,
    timing_cache,
    enable_tf32
)

Build the kwargs for the CreateConfig constructor.

Parameters:

max_workspace_size

Maximum workspace size

optimization_level

TensorRT optimization level

compatibility_level

Hardware compatibility level

profiles

List of optimization profiles

timing_cache

Path to timing cache

enable_tf32

Enable TF32 hardware acceleration

Returns:

Dictionary of kwargs for CreateConfig

aitune.torch.backend.tensorrt.tensorrt_builder.TensorRTBuilder._build_engine_from_network(
    network: typing.Any,
    config: typing.Any
) -> typing.Any

Build TensorRT engine from network.

Parameters:

network

Any

TensorRT network

config

Any

TensorRT configuration

Returns: Any

TensorRT engine

aitune.torch.backend.tensorrt.tensorrt_builder.TensorRTBuilder._create_trt_config(
    profiles: list[polygraphy.backend.trt.Profile]
) -> typing.Any

Create TensorRT config.

Parameters:

profiles

list[Profile]

List of TensorRT optimization profiles

Returns: Any

TensorRT configuration object

aitune.torch.backend.tensorrt.tensorrt_builder.TensorRTBuilder._create_trt_network() -> typing.Any

Create TensorRT network from ONNX.

Returns: Any

TensorRT network

aitune.torch.backend.tensorrt.tensorrt_builder.TensorRTBuilder._handle_failed_build(
    error: Exception,
    engine_path: pathlib.Path
) -> None

Handle errors during engine build.

Parameters:

error

Exception

The exception that occurred

engine_path

Path

Path to the engine file

aitune.torch.backend.tensorrt.tensorrt_builder.TensorRTBuilder._save_engine(
    engine,
    path: pathlib.Path
)

Save the TensorRT engine to a file.

Parameters:

engine

TensorRT engine

path

Path

Path to save the engine

aitune.torch.backend.tensorrt.tensorrt_builder.TensorRTBuilder._validate_onnx_file() -> None

Validate that the ONNX file exists.

Raises:

FileNotFoundError: If the ONNX file does not exist

aitune.torch.backend.tensorrt.tensorrt_builder.TensorRTBuilder.build() -> pathlib.Path

Build the TensorRT engine.

Returns: Path

Path to the TensorRT engine

aitune.torch.backend.tensorrt.tensorrt_builder.logger = logging.getLogger(__name__)

aitune.torch.backend.tensorrt.tensorrt_builder.trt = lazy_import('tensorrt')