aitune.torch.libs.onnx.runtime | NVIDIA AITune Documentation

Shared ONNX Runtime execution with Torch inputs.

Module Contents

Functions

Name	Description
`_bind_inputs`	Bind contiguous CPU or CUDA tensors directly by their memory pointers.
`_bind_outputs`	Tell ORT to allocate all outputs on the requested device.
`_collect_outputs`	Collect CPU outputs as arrays or copy CUDA outputs directly to Torch.
`prepare_onnx_inputs`	Match positional and named inputs to the original ONNX graph names.
`run_onnx`	Bind tensors directly and return outputs on the requested device.

Data

_TORCH_DTYPE_TO_NUMPY

logger

API

aitune.torch.libs.onnx.runtime._bind_inputs(
    io_binding: onnxruntime.IOBinding,
    inputs: dict[str, torch.Tensor]
) -> None

Bind contiguous CPU or CUDA tensors directly by their memory pointers.

aitune.torch.libs.onnx.runtime._bind_outputs(
    io_binding: onnxruntime.IOBinding,
    session: onnxruntime.InferenceSession,
    device: torch.device
) -> None

Tell ORT to allocate all outputs on the requested device.

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

aitune.torch.libs.onnx.runtime._collect_outputs(
    io_binding: onnxruntime.IOBinding,
    session: onnxruntime.InferenceSession,
    device: torch.device
) -> dict[str, torch.Tensor]

Collect CPU outputs as arrays or copy CUDA outputs directly to Torch.

aitune.torch.libs.onnx.runtime.prepare_onnx_inputs(
    session: onnxruntime.InferenceSession,
    args: tuple,
    kwargs: dict
) -> dict[str, torch.Tensor]

Match positional and named inputs to the original ONNX graph names.

aitune.torch.libs.onnx.runtime.run_onnx(
    session: onnxruntime.InferenceSession,
    inputs: dict[str, torch.Tensor],
    device: torch.device
) -> dict[str, torch.Tensor]

Bind tensors directly and return outputs on the requested device.

aitune.torch.libs.onnx.runtime._TORCH_DTYPE_TO_NUMPY: dict[dtype, type] = {torch.float16: np.float16, torch.float32: np.float32, torch.float64: np.float64...

aitune.torch.libs.onnx.runtime.logger = getLogger(__name__)