aitune.torch.backend.tensorrt.tensorrt_backend

TensorRT backend.

Module Contents

Classes

Name	Description
`ProfileMode`	Mode how TRT optimization profiles will be generated for TensorRT engine.
`TensorRTBackend`	TensorRT backend for model acceleration.
`TensorRTBackendConfig`	Configuration for TensorRT backend.
`TensorRTBuildStep`	Identifiers for discrete sub-steps of a TensorRT backend build.
`TensorRTRunner`	TensorRT runner for model acceleration.

Data

ONNX_FILE_EXTENSION

TRT_ENGINE_FILE_EXTENSION

logger

API

class aitune.torch.backend.tensorrt.tensorrt_backend.ProfileMode

Bases: enum.Enum

Mode how TRT optimization profiles will be generated for TensorRT engine.

SAMPLES_USED

= 'samples_used'

SINGLE

= 'single'

class aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend(
    config: aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackendConfig | None = None
)

Bases: Backend, TensorRTRunner

TensorRT backend for model acceleration.

This class provides functionality to build and run TensorRT engines from PyTorch models. It handles the process of exporting models to ONNX and then converting them to TensorRT engines for optimized inference.

STATE_CONFIG

= 'config'

STATE_DEVICE

= 'device'

STATE_ENGINE_PATH

= 'engine_path'

STATE_GRAPH_SPEC

= 'graph_spec'

STATE_OUTPUT_OBJECT

= 'output_object'

STATE_QUANTIZATION_CONFIG

= 'quantization_config'

STATE_TRT_OPTIMIZATION_PROFILES_PATH

= 'trt_optimization_profiles_path'

STATE_TYPE

= 'type'

STATE_USE_CUDA_GRAPHS

= 'use_cuda_graphs'

_build_mode

= BuildMode.AHEAD_OF_TIME

_config

= config or TensorRTBackendConfig()

_cuda_graphs

= TensorRTCudaGraphCache()

_devices

list[str] = ['cuda']

_engine_artifact

ArtifactPath | None = None

_execution_modes

= frozenset({ExecutionMode.SINGLE_GPU})

_trt_optimization_profiles

list[Profile] = []

_trt_optimization_profiles_artifact

ArtifactPath | None = None

_use_cuda_graphs

bool

Use graphs only for fixed-shape profiles, until capture fails for this instance.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._activate()

Activate the TensorRT engine.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._build(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    samples: aitune.torch.module.sample_store.SampleStore,
    cache_dir: pathlib.Path
) -> aitune.torch.backend.backend.Backend

Build the TensorRT model.

This method will export the module to ONNX and build a TensorRT engine from it using Polygraphy.

Parameters:

module

nn.Module

The module to build the TensorRT model for.

name

str

The name of the model.

graph_spec

GraphSpec

The graph spec of the model.

samples

SampleStore

Recorded samples for the model.

cache_dir

Path

The cache directory to store the TensorRT model.

Returns: Backend

The backend with built TensorRT model.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._build_modelopt_onnx(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    samples: collections.abc.Sequence[aitune.torch.module.sample_store.Sample],
    cache_dir: pathlib.Path
) -> pathlib.Path

Build the TensorRT model ModelOpt ONNX quantization.

This method will export the module to ONNX, quantize the exported model using ModelOpt ONNX quantization and build a TensorRT engine from it using Polygraphy.

Parameters:

module

nn.Module

The module to build the TensorRT model for.

name

str

The name of the model.

graph_spec

GraphSpec

The graph spec of the model.

samples

Sequence[Sample]

Recorded samples for the model.

cache_dir

Path

The cache directory to store the TensorRT model.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._build_modelopt_onnx_autocast(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    samples: collections.abc.Sequence[aitune.torch.module.sample_store.Sample],
    cache_dir: pathlib.Path
) -> pathlib.Path

Build the TensorRT model.

This method will export the module to ONNX and build a TensorRT engine from it using Polygraphy.

Parameters:

module

nn.Module

The module to build the TensorRT model for.

name

str

The name of the model.

graph_spec

GraphSpec

The graph spec of the model.

samples

Sequence[Sample]

Recorded samples for the model.

cache_dir

Path

The cache directory to store the TensorRT model.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._build_modelopt_torch(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    samples: collections.abc.Sequence[aitune.torch.module.sample_store.Sample],
    cache_dir: pathlib.Path
) -> pathlib.Path

Build the TensorRT model ModelOpt torch quantization.

This method will quantize module using ModelOpt torch quantization, export the quantized model to ONNX and build a TensorRT engine from it using Polygraphy.

Parameters:

module

nn.Module

The module to build the TensorRT model for.

name

str

The name of the model.

graph_spec

GraphSpec

The graph spec of the model.

samples

Sequence[Sample]

Recorded samples for the model.

cache_dir

Path

The cache directory to store the TensorRT model.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._build_standard(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    samples: collections.abc.Sequence[aitune.torch.module.sample_store.Sample],
    cache_dir: pathlib.Path
) -> pathlib.Path

Build the TensorRT model.

This method will export the module to ONNX and build a TensorRT engine from it using Polygraphy.

Parameters:

module

nn.Module

The module to build the TensorRT model for.

name

str

The name of the model.

graph_spec

GraphSpec

The graph spec of the model.

samples

Sequence[Sample]

Recorded samples for the model.

cache_dir

Path

The cache directory to store the TensorRT model.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._create_output_allocator(
    context
)

Attach an independent output allocator to an execution context.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._deactivate()

Deactivate the TensorRT engine.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._deploy()

Deploys the backend.

After deploying, the backend is ready to do inference. Backend cannot be deactivated anymore.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._execute_engine()

Execute TensorRT normally, propagating engine execution failures.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._find_optimization_profile(
    inputs: dict[str, torch.Tensor]
) -> int

Find the first matching profile, including validation for a single profile.

Parameters:

inputs

dict[str, torch.Tensor]

Dictionary mapping input names to tensors

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._get_output_object(
    module: torch.nn.Module,
    sample: aitune.torch.module.sample_store.Sample
) -> typing.Any

Get the output object from the module and sample.

Note: to avoid case where a module returns a reference to the input argument, we make a deep copy of the output object.

Parameters:

module

nn.Module

PyTorch module

sample

Sample

Sample input to use for model inference.

Returns: Any

The output object from the module.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._get_profiles_from_shapes() -> list[polygraphy.backend.trt.Profile]

Create TensorRT optimization profiles.

Returns: list[Profile]

List of Polygraphy Profile objects

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._get_shapes(
    graph_spec: aitune.torch.module.graph_spec.GraphSpec
) -> tuple[dict[str, tuple[int, ...]], dict[str, tuple[int, ...]], dict[str, tuple[int, ...]]]

Get the shapes for the input tensors.

Parameters:

graph_spec

GraphSpec

Input graph spec

Returns: tuple[dict[str, tuple[int, ...]], dict[str, tuple[int, ...]], dict[str, tuple[int, ...]]]

Tuple of dictionaries of minimum, optimal, and maximum input shapes

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._infer(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Infer using the TensorRT engine.

Parameters:

*args

AnyDefaults to ()

Input tensors

**kwargs

AnyDefaults to {}

Named input tensors

Returns: Any

Model outputs

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._load_trt_optimization_profiles(
    trt_optimization_profiles_path: pathlib.Path
) -> list[polygraphy.backend.trt.Profile]

Load the TensorRT optimization profiles from a file.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._prepare_inputs(
    args,
    kwargs
)

Prepare input tensors from args and kwargs.

Parameters:

args

Positional input tensors

kwargs

Named input tensors

Returns:

Dictionary mapping input names to tensors

Raises:

ValueError: If inputs are missing or incorrect

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._prepare_onnx_model_path(
    cache_dir: pathlib.Path,
    suffix: str = ''
) -> pathlib.Path

Prepare the ONNX model path.

Parameters:

cache_dir

Path

The cache directory to store the ONNX model.

suffix

strDefaults to ''

The suffix of the ONNX model.

Returns: Path

The ONNX model path

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._prepare_outputs_for_return() -> typing.Any

Prepare the outputs for return.

This method will prepare the outputs for return according to the original model’s output structure. Tensors are retrieved from the output allocator with correct shapes.

Returns: Any

The outputs in the same format as the original model.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._prepare_trt_engine_path(
    cache_dir: pathlib.Path
) -> pathlib.Path

Prepare the TensorRT engine path.

Parameters:

cache_dir

Path

The cache directory to store the TensorRT model.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._prepare_trt_optimization_profiles_path(
    cache_dir: pathlib.Path
) -> pathlib.Path

Prepare the TensorRT optimization profiles path.

Parameters:

cache_dir

Path

The cache directory to store the TensorRT optimization profiles.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._save_trt_optimization_profiles(
    optimization_profiles: list[polygraphy.backend.trt.Profile],
    cache_dir: pathlib.Path
) -> pathlib.Path

Save the TensorRT optimization profiles to a file.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._set_input_tensors(
    inputs
)

Set shapes and memory addresses for input tensors.

Parameters:

inputs

Dictionary mapping input names to tensors

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._set_optimization_profiles(
    inputs: dict[str, torch.Tensor]
) -> None

Select an ordinary context or the static profile’s cached graph context.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._timing_cache_path() -> pathlib.Path | None

Return a timing cache path that is safe for this process.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend.describe() -> str

Returns the description of the backend.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend.from_dict(
    module: torch.nn.Module,
    state_dict: dict
)

classmethod

Creates a backend from a state_dict.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend.get_profiles(
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    samples: collections.abc.Sequence[aitune.torch.module.sample_store.Sample]
) -> list[polygraphy.backend.trt.Profile]

Create profiles from samples or from graph_spec.

If self._config.profiles is a list, return the user provided profiles. If self._config.profiles is ProfileMode.SINGLE, create a single profile from the graph spec. If self._config.profiles is ProfileMode.SAMPLES_USED, create profiles from shapes seen in samples. Explicit module shape definitions always produce a single authoritative profile and cannot be combined with user-provided TensorRT profiles.

Parameters:

graph_spec

GraphSpec

Input graph spec

samples

Sequence[Sample]

Recorded samples for the model.

Returns: list[Profile]

List of The Polygraphy Profile objects

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend.key() -> str

Returns the key of the backend.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend.to_dict()

Returns the state_dict of the backend.

class aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackendConfig(
    use_dynamo: bool = True,
    workspace_size: int | None = None,
    opset_version: int | None = None,
    optimization_level: int | None = None,
    compatibility_level: int | None = None,
    timing_cache: pathlib.Path | None = None,
    profiles: aitune.torch.backend.tensorrt.tensorrt_backend.ProfileMode | list[aitune.torch.backend.tensorrt.tensorrt_profile.TensorRTProfile] = ProfileMode.SINGLE,
    device: str = 'cuda',
    quantization_config: aitune.torch.backend.tensorrt.onnx_autocast.ONNXAutoCastConfig | aitune.torch.backend.tensorrt.onnx_quantization.ONNXQuantizationConfig | aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizationConfig | None = None,
    enable_tf32: bool = True,
    use_cuda_graphs: bool = True,
    max_cuda_graphs: int = 8,
    cuda_graph_cache_policy: aitune.torch.backend.tensorrt.cuda_graphs.CudaGraphCachePolicy = 'lfu'
)

Dataclass

Bases: BackendConfig

Configuration for TensorRT backend.

compatibility_level

int | None = None

cuda_graph_cache_policy

CudaGraphCachePolicy = 'lfu'

device

str = 'cuda'

enable_tf32

bool = True

max_cuda_graphs

int = 8

opset_version

int | None = None

optimization_level

int | None = None

profiles

ProfileMode | list[TensorRTProfile] = ProfileMode.SINGLE

quantization_config

ONNXAutoCastConfig | ONNXQuantizationConfig | TorchQuantizationConfig | None = None

timing_cache

Path | None = None

use_cuda_graphs

bool = True

use_dynamo

bool = True

workspace_size

int | None = None

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackendConfig.__post_init__()

Validate the graph cache capacity and policy.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackendConfig.from_dict(
    data: dict
) -> aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackendConfig

classmethod

Initialise config from a plain dict (e.g. parsed from YAML).

profiles may be passed as a string (ProfileMode value) or a list of profile dicts and will be reconstructed automatically. quantization_config may be passed as a dict with a _type key (produced by to_dict()) and will be reconstructed automatically.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackendConfig.profiles_from_dict(
    data: str | list[dict]
) -> aitune.torch.backend.tensorrt.tensorrt_backend.ProfileMode | list[aitune.torch.backend.tensorrt.tensorrt_profile.TensorRTProfile]

classmethod

Convert dict to list of TensorRTProfile.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackendConfig.quantization_config_from_dict(
    data: dict
) -> aitune.torch.backend.tensorrt.onnx_autocast.ONNXAutoCastConfig | aitune.torch.backend.tensorrt.onnx_quantization.ONNXQuantizationConfig | aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizationConfig

classmethod

Reconstruct a quantization config from a dict produced by to_dict().

The dict must contain a _type key with the class name.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackendConfig.to_dict() -> dict

Convert TensorRTBackendConfig to dictionary.

class aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBuildStep()

Bases: BackendBuildStep

Identifiers for discrete sub-steps of a TensorRT backend build.

MODELOPT_ONNX_QUANTIZATION

= 'ModelOpt ONNX quantization'

MODELOPT_TORCH_QUANTIZATION

= 'ModelOpt Torch quantization'

ONNX_AUTOCAST

= 'ONNX autocast'

ONNX_EXPORT

= 'ONNX export'

TENSORRT_ENGINE_BUILD

= 'TensorRT engine build'

class aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTRunner(
    args = (),
    kwargs = {}
)

TensorRT runner for model acceleration.

This class provides functionality to run TensorRT engines from PyTorch models.

aitune.torch.backend.tensorrt.tensorrt_backend.ONNX_FILE_EXTENSION = '.onnx'

aitune.torch.backend.tensorrt.tensorrt_backend.TRT_ENGINE_FILE_EXTENSION = '.plan'

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