onnx_utils - Infernet ML
Description: Infernet ML Documentation

Skip to content

# Module: `onnx_utils`

## `generate_dummy_input(model_path)`

Generate dummy input for an ONNX model, read the input shape from the model and generate random input data.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`model_path` | `Path` | Path to the ONNX model | required

Returns:

Name | Type | Description
--- | --- | ---
`ONNXInput` | `ONNXInput` | Dictionary containing dummy input for the model. Keys are the input names and values are the dummy input data.

Source code in `src/infernet_ml/utils/onnx_utils.py`

     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    def generate_dummy_input(model_path: Path) -> ONNXInput:
        """
        Generate dummy input for an ONNX model, read the input shape from the model
        and generate random input data.

        Args:
            model_path (Path): Path to the ONNX model

        Returns:
            ONNXInput: Dictionary containing dummy input for the model. Keys are the
                input names and values are the dummy input data.
        """

        # Load the ONNX model
        model = onnx.load(model_path)

        # Initialize a dictionary to store dummy inputs
        dummy_inputs = {}

        # Iterate over each input in the model's graph
        for input_tensor in model.graph.input:
            # Get the shape of the input
            shape = []
            for dim in input_tensor.type.tensor_type.shape.dim:
                # Handle cases where dimension is not defined
                if dim.dim_value > 0:
                    shape.append(dim.dim_value)
                else:
                    # If dimension value is not set, use a default value (e.g., 1)
                    shape.append(1)

            # Determine the data type
            dtype = onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[
                input_tensor.type.tensor_type.elem_type
            ]

            # Generate a random input tensor with the specified shape and data type
            dummy_input = np.random.random_sample(shape).astype(dtype)

            # Store the dummy input in the dictionary
            dummy_inputs[input_tensor.name] = dummy_input

        return dummy_inputs

## `get_onnx_input_names(onnx_file_path)`

Get the input names of an ONNX model, from the model file.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`onnx_file_path` | `Path` | Path to the ONNX model file | required

Returns:

Type | Description
--- | ---
`List[str]` | List[str]: List of input names

Source code in `src/infernet_ml/utils/onnx_utils.py`

    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    def get_onnx_input_names(onnx_file_path: Path) -> List[str]:
        """
        Get the input names of an ONNX model, from the model file.

        Args:
            onnx_file_path (Path): Path to the ONNX model file

        Returns:
            List[str]: List of input names
        """
        # Load the ONNX model
        model = onnx.load(onnx_file_path)

        # Get the model's graph
        graph = model.graph

        # Get the inputs of the graph
        inputs = graph.input

        # Check if there are any inputs
        if not inputs:
            raise ValueError("The model has no input feeds.")

        # Collect all input names
        input_names = [input.name for input in inputs]

        return input_names

## `get_onnx_output_names(onnx_file_path)`

Get the output names of an ONNX model, from the model file.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`onnx_file_path` | `Path` | Path to the ONNX model file | required

Returns:

Type | Description
--- | ---
`List[str]` | List[str]: List of output names

Source code in `src/infernet_ml/utils/onnx_utils.py`

    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    def get_onnx_output_names(onnx_file_path: Path) -> List[str]:
        """
        Get the output names of an ONNX model, from the model file.

        Args:
            onnx_file_path (Path): Path to the ONNX model file

        Returns:
            List[str]: List of output names
        """
        # Load the ONNX model
        model = onnx.load(onnx_file_path)

        # Get the model's graph
        graph = model.graph

        # Get the outputs of the graph
        outputs = graph.output

        # Check if there are any outputs
        if not outputs:
            raise ValueError("The model has no output nodes.")

        # Collect all output names
        output_names = [output.name for output in outputs]

        return output_names