ONNX Runtime backend.

Module Contents

Classes

Name	Description
`ONNXExecutionProvider`	Supported ONNX Runtime execution providers.
`ONNXRuntimeBackend`	Backend that exports models to ONNX and runs inference with ONNX Runtime.
`ONNXRuntimeBackendConfig`	Configuration for ONNXRuntimeBackend.

Data

_TORCH_DTYPE_TO_NUMPY

logger

API

class aitune.torch.backend.onnx_runtime_backend.ONNXExecutionProvider

Bases: enum.Enum

Supported ONNX Runtime execution providers.

Only NVIDIA GPU-backed providers are supported:

CUDA — CUDAExecutionProvider: standard GPU execution.
TENSORRT — TensorrtExecutionProvider with CUDAExecutionProvider as fallback: enables TensorRT engine compilation for maximum GPU throughput.

CUDA

= 'cuda'

TENSORRT

= 'tensorrt'

class aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend(
    config: aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackendConfig | None = None
)

Bases: Backend

Backend that exports models to ONNX and runs inference with ONNX Runtime.

Exports the model to a .onnx artifact at build time (trace or dynamo), then loads an onnxruntime.InferenceSession for inference. Dynamic batch and spatial dimensions are inferred automatically from graph_spec.

Workflow::

backend = ONNXRuntimeBackend()

build / tune as usual through ait.tune()

ait.save(model, “model.ait”)

later

ait.load(model, “model.ait”)

STATE_CONFIG

= 'config'

STATE_DEVICE

= 'device'

STATE_GRAPH_SPEC

= 'graph_spec'

STATE_ONNX_DATA_PATH

= 'onnx_data_path'

STATE_ONNX_MODEL_PATH

= 'onnx_model_path'

STATE_OUTPUT_OBJECT

= 'output_object'

STATE_TYPE

= 'type'

_config

= config or ONNXRuntimeBackendConfig()

_devices

list[str] = ['cuda']

_graph_spec

GraphSpec | None = None

_onnx_data_path

Path | None = None

_onnx_model_path

Path | None = None

_session

InferenceSession | None = None

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend._activate()

Load the ONNX Runtime session from disk.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend._bind_inputs(
    io_binding: onnxruntime.IOBinding,
    inputs: dict[str, typing.Any]
) -> None

Bind prepared inputs to an IOBinding handle.

GPU tensors are bound zero-copy via DLPack; CPU tensors and other array-like values are bound as numpy arrays.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend._bind_outputs(
    io_binding: onnxruntime.IOBinding
) -> None

Tell ORT to allocate all outputs on the CUDA device.

ORT owns the output buffers; shapes are resolved at inference time. Tensors are retrieved after inference via _collect_outputs.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend._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

Export the model to ONNX then load the session.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend._collect_outputs(
    io_binding: onnxruntime.IOBinding
) -> dict[str, torch.Tensor]

Collect ORT CUDA outputs into torch tensors via D2D memcpy (no CPU round-trip).

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend._deactivate()

Deactivate backend.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend._deploy()

Deploy backend.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend._get_execution_providers() -> list[str | tuple]

Return execution providers based on config (no automatic CPU fallback).

CUDA → [("CUDAExecutionProvider", {"device_id": ...})]
TENSORRT → ["TensorrtExecutionProvider", ("CUDAExecutionProvider", {"device_id": ...})]

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend._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.onnx_runtime_backend.ONNXRuntimeBackend._infer(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Run inference through the ONNX Runtime session via IOBinding.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend._prepare_inputs(
    args: tuple,
    kwargs: dict
) -> dict[str, typing.Any]

Map args/kwargs to session input names using graph_spec locators.

Tensors are returned as-is (preserving their device); conversion to the format expected by ONNX Runtime happens in _infer via IOBinding.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend._prepare_outputs(
    outputs: dict[str, torch.Tensor]
) -> typing.Any

Reconstruct original output structure from session output tensors.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend._save_config(
    cache_dir: pathlib.Path
)

Store the backend configuration to a file.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend.describe() -> str

Returns the description of the backend.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend.from_dict(
    module: torch.nn.Module | None,
    state_dict: dict
) -> aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend

classmethod

Creates a backend from a state_dict.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend.key() -> str

Returns the key of the backend.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackend.to_dict() -> dict

Returns the state_dict of the backend.

class aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackendConfig(
    use_dynamo: bool = True,
    execution_provider: aitune.torch.backend.onnx_runtime_backend.ONNXExecutionProvider | None = None,
    opset_version: int | None = None
)

Dataclass

Bases: BackendConfig

Configuration for ONNXRuntimeBackend.

Parameters:

use_dynamo

boolDefaults to True

If True (default), export via torch.onnx.export(dynamo=True) which calls torch.export.export internally and produces a more accurate graph (no Python-level tracing limitations). If False, use the classic trace-based exporter — faster and broader model coverage.

execution_provider

ONNXExecutionProvider | NoneDefaults to None

ONNX Runtime execution provider to use. When None the backend defaults to :attr:ONNXExecutionProvider.CUDA. Only :attr:ONNXExecutionProvider.CUDA and :attr:ONNXExecutionProvider.TENSORRT are supported.

opset_version

int | NoneDefaults to None

ONNX opset version passed to torch.onnx.export. None uses the torch default.

execution_provider

ONNXExecutionProvider | None = None

opset_version

int | None = None

use_dynamo

bool = True

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackendConfig.__post_init__()

Post init.

aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackendConfig.from_dict(
    data: dict
) -> aitune.torch.backend.onnx_runtime_backend.ONNXRuntimeBackendConfig

classmethod

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

execution_provider may be passed as a string and will be converted to an ONNXExecutionProvider enum value automatically.

aitune.torch.backend.onnx_runtime_backend._TORCH_DTYPE_TO_NUMPY: dict[dtype, type] = {torch.float16: np.float16, torch.float32: np.float32, torch.float64: np.float64...

aitune.torch.backend.onnx_runtime_backend.logger = getLogger(__name__)