OnnxRuntime#

API#

class warp_nn.runtime.OnnxRuntime(

path: str,

*,

device: str | Device | None = None,

requires_grad: bool = False,

**kwargs: Any,

)[source]#

Bases: object

Execute an ONNX graph using Warp-NN modules.

Each graph node is compiled into Warp-NN modules (e.g. Gemm into Linear, LSTM into LSTM), with its constant inputs (initializers and Constant nodes) loaded into the modules’ parameters and buffers.

Symbolic ONNX dimensions are resolved from the actual input arrays. The first call with a given set of input shapes allocates the (cached) output arrays of the modules and loads their kernels; subsequent calls with the same input shapes are allocation-free and can be captured in a CUDA graph (see prepare()). The returned arrays are owned by the runtime and overwritten by subsequent calls (or, for pass-through nodes such as Identity, are the input arrays themselves), so they must not be written to.

Parameters:

path – Path to an ONNX model.
device – Target Warp device. None selects the current default device.
requires_grad – Whether the module parameters and cached output arrays require gradients.

Raises:

NotImplementedError – If the model uses unsupported operators, attributes or inputs.
ValueError – If the model is invalid.

Example:

>>> import warp as wp
>>> from warp_nn.runtime import OnnxRuntime
>>>
>>> # load the ONNX model
>>> runtime = OnnxRuntime("policy.onnx")
>>>
>>> # run the graph with a batch of 4 observations
>>> outputs = runtime({"obs": wp.zeros((4, 6), dtype=wp.float32)})
>>> outputs["actions"].shape
(4, 3)

__call__( inputs: dict[str, array], ) → dict[str, array][source]#

Execute the graph.

Parameters:: inputs – Input arrays, keyed by graph input name. They must be contiguous.
Returns:: Output arrays, keyed by graph output name.

Example:

>>> obs = wp.zeros((4, 6), dtype=wp.float32)
>>> outputs = runtime({"obs": obs})
>>> list(outputs)
['actions']
>>> runtime({"obs": obs})["actions"].ptr == outputs["actions"].ptr  # the outputs are cached per input shape
True

prepare( inputs: dict[str, array] | None = None, *, batch_size: int | None = None, ) → None[source]#

Run the graph once to allocate the cached arrays and load the kernels for the given input shapes.

This must be done before capturing the runtime calls in a CUDA graph, since neither memory allocations nor kernel loading can be captured.

Exactly one of inputs or batch_size must be given. When batch_size is given, one-filled inputs are generated from the ONNX declarations, with the symbolic (or unknown) batch dimension set to batch_size. Each input must then declare at most one symbolic dimension, and all the named symbolic dimensions must share the same name; otherwise, the inputs must be given explicitly. Note that the symbolic dimension is assumed to be the batch one, so models with a fixed batch dimension and another symbolic dimension (e.g. a sequence length) must also be prepared with explicit inputs.

Parameters:

inputs – Input arrays, keyed by graph input name.
batch_size – Size of the batch dimension.

Raises:

ValueError – If both or neither of inputs and batch_size are given, if batch_size is not positive, or if the batch dimension is ambiguous.

Example:

>>> from warp_nn.utils import ScopedCapture
>>>
>>> # prepare the runtime to allow capturing the calls in a CUDA graph
>>> runtime.prepare(batch_size=4)  # equivalent to runtime.prepare({"obs": obs})
>>>
>>> obs = wp.zeros((4, 6), dtype=wp.float32)
>>> with ScopedCapture(enabled=wp.get_device().is_cuda) as capture:  # capture on CUDA devices only
...     actions = runtime({"obs": obs})["actions"]
>>>
>>> # launch the captured graph or run the runtime normally if no graph was captured
>>> obs.fill_(1.0)  # update the inputs in place
>>> if capture.graph is not None:
...     wp.capture_launch(capture.graph)
... else:
...     actions = runtime({"obs": obs})["actions"]
>>> actions.shape
(4, 3)

property _shapes: dict[str, tuple[int, ...]][source]#

Shapes of the graph inputs and outputs, keyed by name.

The shapes are those of the last prepare() call, or the declared ones (with the symbolic dimensions set to 1) if the runtime has not been prepared.

Deprecated since version 0.4.0: Use inputs/outputs, or the shapes of the arrays returned by the runtime, instead.

property input_names: list[str][source]#: Names of the graph inputs (excluding initializers), in declaration order.

Deprecated since version 0.4.0: Use inputs instead.

property inputs: tuple[OnnxTensorSpec, ...][source]#

Name, shape and data type of the graph inputs (excluding initializers), in declaration order.

Example:

>>> runtime.inputs
(OnnxTensorSpec(name='obs', shape=(None, 6), dtype=<class '...float32'>),)

property output_names: list[str][source]#: Names of the graph outputs, in declaration order.

Deprecated since version 0.4.0: Use outputs instead.

property outputs: tuple[OnnxTensorSpec, ...][source]#

Name, shape and data type of the graph outputs, in declaration order.

Example:

>>> runtime.outputs
(OnnxTensorSpec(name='actions', shape=(None, 3), dtype=<class '...float32'>),)

class warp_nn.runtime.OnnxTensorSpec( *, name: str, shape: tuple[int | None, ...], dtype: type, )[source]#

Bases: object

Name, shape and data type of an ONNX graph input or output.

Symbolic (and otherwise unknown) dimensions are represented by None.

classmethod from_onnx( value: Any, ) → OnnxTensorSpec[source]#: Create a tensor specification, preserving symbolic dimensions.

dtype: type#

name: str#

shape: tuple[int | None, ...]#