aitune.torch.module.onnx_module | NVIDIA AITune Documentation

A runnable ONNX graph for AITune recording and tuning.

Module Contents

Classes

Name	Description
`OnnxModule`	Run an ONNX file with positional or named tensor inputs and named outputs.

Functions

Name	Description
`onnx_checkpoint_placeholder`	Supply a source module when loading a self-contained ONNX AOT checkpoint.

API

class aitune.torch.module.OnnxModule(
    path: str | pathlib.Path
)

Bases: Module

Run an ONNX file with positional or named tensor inputs and named outputs.

Wrap with aitune.torch.Module to record calls and select a backend with MaxThroughputStrategy. Inputs must match the ONNX graph’s names and dtypes.

_input_names

tuple[str, ...] = ()

_output_names

tuple[str, ...] = ()

graph_info

ONNXModelInfo

Read metadata declared by the source ONNX graph without running inference.

path

= Path(path).resolve()

aitune.torch.module.OnnxModule._ensure_session(
    device: torch.device
) -> None

Create a session lazily without switching an existing session’s device.

aitune.torch.module.OnnxModule._providers(
    device: torch.device
) -> list

staticmethod

Return execution providers for the selected device.

aitune.torch.module.OnnxModule.deactivate() -> None

Release the runtime session.

aitune.torch.module.OnnxModule.forward(
    args: torch.Tensor = (),
    kwargs: torch.Tensor = {}
) -> dict[str, torch.Tensor]

Execute the graph using the same runtime as ONNXRuntimeBackend.

aitune.torch.module.OnnxModule.offload(
    device: str | torch.device = 'cpu'
) -> None

Switch execution providers, using CPU for meta since ORT has no meta provider.

aitune.torch.module.OnnxModule.preserve_tensor_names(
    graph_spec: aitune.torch.module.graph_spec.GraphSpec
) -> None

Attach original ONNX names to recorded tensors without altering their observed shapes.

aitune.torch.module.onnx_checkpoint_placeholder() -> torch.nn.Identity

Supply a source module when loading a self-contained ONNX AOT checkpoint.