> This page is for AITune, version v0.4.1.
> For other versions, use one of these documentation indexes:
> - Latest (v0.6.0) (default): https://docs.nvidia.com/aitune/latest/llms.txt
> - dev: https://docs.nvidia.com/aitune/dev/llms.txt
> - v0.6.0: https://docs.nvidia.com/aitune/v0.6.0/llms.txt
> - v0.5.0: https://docs.nvidia.com/aitune/v0.5.0/llms.txt
> - v0.4.1: https://docs.nvidia.com/aitune/v0.4.1/llms.txt

> For clean Markdown content of this page, append .md to this URL.

# aitune.torch.backend.tensorrt.tensorrt_backend

TensorRT backend.

## Module Contents

### Classes

| Name                                                                                             | Description                                                               |
| ------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------- |
| [`ProfileMode`](#aitune-torch-backend-tensorrt-tensorrt_backend-ProfileMode)                     | Mode how TRT optimization profiles will be generated for TensorRT engine. |
| [`TensorRTBackend`](#aitune-torch-backend-tensorrt-tensorrt_backend-TensorRTBackend)             | TensorRT backend for model acceleration.                                  |
| [`TensorRTBackendConfig`](#aitune-torch-backend-tensorrt-tensorrt_backend-TensorRTBackendConfig) | Configuration for TensorRT backend.                                       |
| [`TensorRTBuildStep`](#aitune-torch-backend-tensorrt-tensorrt_backend-TensorRTBuildStep)         | Identifiers for discrete sub-steps of a TensorRT backend build.           |
| [`TensorRTRunner`](#aitune-torch-backend-tensorrt-tensorrt_backend-TensorRTRunner)               | TensorRT runner for model acceleration.                                   |

### Data

[`ONNX_FILE_EXTENSION`](#aitune-torch-backend-tensorrt-tensorrt_backend-ONNX_FILE_EXTENSION)

[`TRT_ENGINE_FILE_EXTENSION`](#aitune-torch-backend-tensorrt-tensorrt_backend-TRT_ENGINE_FILE_EXTENSION)

[`logger`](#aitune-torch-backend-tensorrt-tensorrt_backend-logger)

### API

```python
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'`

---

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

**Bases:** [Backend](/aitune/aitune/torch/backend/backend#aitune-torch-backend-backend-Backend), [TensorRTRunner](#aitune-torch-backend-tensorrt-tensorrt_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] = []`

---

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

Activate the TensorRT engine.

```python
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.

```python
aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._build_cuda_graph()
```

Create a CUDA graph for inference.

```python
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.

---

```python
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.

---

```python
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.

---

```python
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.

---

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

Deactivate the TensorRT engine.

```python
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.

```python
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:

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

```python
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.

```python
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

```python
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

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

Infer using the TensorRT engine.

**Parameters:**

**`*args`** `Any` — default: ()

Input tensors

---

**`**kwargs`** `Any` — default: \{}

Named input tensors

---

**Returns:** `Any`

Model outputs

```python
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.

```python
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.

```python
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

```python
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`** `str` — default: ''

The suffix of the ONNX model.

---

**Returns:** `Path`

The ONNX model path

```python
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.

```python
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.

---

```python
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.

---

```python
aitune.torch.backend.tensorrt.tensorrt_backend.TensorRTBackend._save_config(
    cache_dir: pathlib.Path
)
```

Store the backend configuration to a file.

```python
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.

```python
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

---

```python
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

---

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

Returns the description of the backend.

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

classmethod

Creates a backend from a state_dict.

```python
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.

Returns:
List of The Polygraphy Profile objects

**Parameters:**

**`graph_spec`** `GraphSpec`

Input graph spec

---

**`data`** `list[Sample]`

List of samples

---

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

Returns the key of the backend.

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

Returns the state_dict of the backend.

```python
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](/aitune/aitune/torch/backend/backend#aitune-torch-backend-backend-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`

---

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

Returns the default fields to describe.

```python
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.

```python
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.

```python
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.

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

Convert TensorRTBackendConfig to dictionary.

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

**Bases:** [BackendBuildStep](/aitune/aitune/torch/backend/backend#aitune-torch-backend-backend-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'`

---

```python
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.

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

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

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