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_SAMPLES

= 'samples'

STATE_TYPE

= 'type'

_build_mode

= BuildMode.AHEAD_OF_TIME

_config

= config or ONNXRuntimeBackendConfig()

_devices

list[str] = ['cuda']

_execution_modes

= frozenset({ExecutionMode.SINGLE_GPU})

_graph_spec

GraphSpec | None = None

_onnx_data_artifact

ArtifactPath | None = None

_onnx_model_artifact

ArtifactPath | None = None

_samples

SampleStore | 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,
    samples: aitune.torch.module.sample_store.SampleStore,
    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.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.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._warmup(
    samples: collections.abc.Iterable[aitune.torch.module.sample_store.Sample]
) -> None

Run representative samples to initialize the execution provider.

Some execution providers defer work until inference. In particular, the TensorRT execution provider compiles the ONNX graph when it first runs. Exercising a small number of recorded samples makes those failures part of the backend build instead of deferring them until profiling or application inference.

Parameters:

samples

Iterable[Sample]

Recorded input samples for the backend build.

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