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'

_config

= config or TensorRTBackendConfig()

_devices

list[str] = ['cuda']

_static_inputs

= {}

_trt_optimization_profiles

list[Profile] = []

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,
    data: list[aitune.torch.module.recording_module.Sample],
    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.

data

list[Sample]

The data of 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_cuda_graph()

Create a CUDA graph for inference.

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._build_modelopt_onnx(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.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.

data

list[Sample]

The data of 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,
    data: list[aitune.torch.module.recording_module.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.

data

list[Sample]

The data of 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,
    data: list[aitune.torch.module.recording_module.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.

data

list[Sample]

The data of 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,
    data: list[aitune.torch.module.recording_module.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.

data

list[Sample]

The data of the model.

cache_dir

Path

The cache directory to store the TensorRT model.

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_cuda_graph()

Execute inference using CUDA graphs for optimized performance.

This method implements the CUDA graph capture and launch pattern:

If no CUDA graph exists, capture one by running inference twice
If CUDA graph exists, launch the captured graph

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._get_output_object(
    module: torch.nn.Module,
    sample: aitune.torch.module.recording_module.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._invalidate_cuda_graph(
    inputs: dict[str, torch.Tensor]
)

Setup the inputs for the CUDA graph.

This should be called when input shapes change or when the graph needs to be rebuilt.

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_config(
    cache_dir: pathlib.Path
)

Store the backend configuration to a file.

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]
)

Set optimization profiles for the input tensors.

Parameters:

inputs

dict[str, torch.Tensor]

Dictionary mapping input names to tensors

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,
    data: list[aitune.torch.module.recording_module.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.

Returns: List of The Polygraphy Profile objects

Parameters:

graph_spec

GraphSpec

Input graph spec

data

list[Sample]

List of samples

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 = False
)

Dataclass

Bases: BackendConfig

Configuration for TensorRT backend.

compatibility_level

int | None = None

device

str = 'cuda'

enable_tf32

bool = True

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 = False

use_dynamo

bool = True

workspace_size

int | None = None

aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackendConfig._default_describe_fields() -> list[str]

Returns the default fields to describe.

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__)