yobx.helpers.onnx_helper

yobx.helpers.onnx_helper.dtype_to_tensor_dtype(dt: dtype | torch.dtype) → int[source][source]

Converts a torch dtype or numpy dtype into a onnx element type.

Parameters:

to – dtype

Returns:

onnx type

yobx.helpers.onnx_helper.get_hidden_inputs(graph: GraphProto) → Set[str][source][source]

Returns the hidden inputs (inputs coming from an upper context) used by a subgraph. It excludes empty names.

yobx.helpers.onnx_helper.np_dtype_to_tensor_dtype(dtype: dtype) → int[source][source]

Converts a numpy dtype to an onnx element type.

yobx.helpers.onnx_helper.onnx_dtype_name(itype: int, exc: bool = True) → str[source][source]

Returns the ONNX name for a specific element type.

<<<

import onnx
from yobx.helpers.onnx_helper import onnx_dtype_name

itype = onnx.onnx.TensorProto.BFLOAT16
print(onnx_dtype_name(itype))
print(onnx_dtype_name(7))

>>>

    BFLOAT16
    INT64
yobx.helpers.onnx_helper.pretty_onnx(onx: AttributeProto | FunctionProto | GraphProto | ModelProto | NodeProto | SparseTensorProto | TensorProto | ValueInfoProto | str, with_attributes: bool = False, highlight: Set[str] | None = None, shape_inference: bool = False) → str[source][source]

Displays an onnx proto in a better way.

Parameters:

  • with_attributes – displays attributes as well, if only a node is printed

  • highlight – to highlight some names

  • shape_inference – run shape inference before printing the model

Returns:

text

yobx.helpers.onnx_helper.tensor_dtype_to_np_dtype(tensor_dtype: int) → dtype[source][source]

Converts a onnx.TensorProto’s data_type to corresponding numpy dtype. It can be used while making tensor.

Parameters:

tensor_dtype – onnx.TensorProto’s data_type

Returns:

numpy’s data_type