Aidge ONNX tutorial

In this tutorial, we will see how to extensivly use the Aidge ONNX module.

The following points will be covered:

• How to load an ONNX;

• How to add the support for an ONNX operator;

• How to support an unsupported operator for export purposes;

• How to add an implementation to an unknown operator.

Install requirements

Ensure that the Aidge modules are properly installed in the current environment. If it is the case, the following setup steps can be skipped.

Note: When running this notebook on Binder, all required components are pre-installed.

%pip install aidge-core \
    aidge-backend-cpu \
    aidge-onnx \
    aidge-model-explorer

Setting up the notebook

Import required modules

import aidge_core
import aidge_onnx
import aidge_backend_cpu # Required for Producer implementation
import numpy as np       # Required to load data

Retrieve the ONNX model

In order to run this tutorial, we will use a simple ONNX composed of a single operator Swish. This operator is not supported by ONNX and is often decomposed in multiple operators.

If you don’t have git-lfs, you can download the model and data using code snippet:

BASE_URL = "https://gitlab.eclipse.org/eclipse/aidge/aidge/-/raw/main/examples/tutorials/Onnx_tutorial/"
file_name= "test_swish.onnx"

# Download the ONNX model file
aidge_core.utils.download_file(file_path=file_name, file_url=f"{BASE_URL}{file_name}")

Importing an ONNX

Importing an ONNX using Aidge is done using the function: aidge_onnx.load_onnx().

graph = aidge_onnx.load_onnx("test_swish.onnx")

The Swish operator is not supported and thus is loaded as a GenericOperator. This mechanism allows to load an unsupported graph into the framework.

The aidge_onnx library has a coverage report tools in order to check how well the graph loaded is supported:

if not aidge_onnx.has_native_coverage(graph):
    print("The graph is not fully supported by Aidge!\n")
aidge_onnx.native_coverage_report(graph)

However, this does not mean we cannot work with this graph!

Working with Generic Operators

Indeed, using the python library, we can work with a GenericOperator.

For this we will begin with retrieving the operator:

swish_node = graph.get_nodes().pop()   # get_nodes() returns a set
swish_op = swish_node.get_operator()   # Retrieving Operator from Node

Computing output dimensions

In order to generate a scheduling, we need to specify how the operator will modify the data. This is required so that Aidge can propagate input/output dimensions.

Generating a scheduling is necessary to generate an export or make inference in the framework.

We can set a function to compute the dimensions using the set_forward_dims() method.

In our case, the Swish function does not modify the dimensions so we can just send the same dimensions as the input. For this, we will create an identity lambda function.

swish_op.set_forward_dims(lambda x: x)

Providing an implementation

To perform inference, we need to provide an implementation; specifically, we must define the forward function.

The Swish function is defined as: \(swish(x)={{x}\over{1+e^{-\beta x}}}\).

Thus, we can create a simple implementation using the Numpy library:

from functools import reduce

class SwishImpl(aidge_core.OperatorImpl): # Inherit OperatorImpl to interface with Aidge !
    def __init__(self, op: aidge_core.Operator):
        aidge_core.OperatorImpl.__init__(self, op, "swish") # Required to avoid type error with C++ binding !
        self.op = op # Reference to the Aidge operator to retrieve attributes, inputs, outputs ..
    def forward(self):
        data_input = np.array(self.op.get_input(0))
        beta = np.array(self.op.attr.get_attr("beta")) # Attribute name is the same as the one in the ONNX
        output =  (data_input / (1 + np.exp(-data_input*beta)))
        self.op.set_output(0, aidge_core.Tensor(output)) # setting operator output

This implementation can then be set using:

swish_op.set_impl(SwishImpl(swish_op)) # Setting implementation

Once this is done, we can run an inference.

Let’s first create an input:

numpy_tensor = np.random.randn(1, 10).astype(np.float32)
in_tensor = aidge_core.Tensor(numpy_tensor)
print(f"Random input:\n{numpy_tensor}")

Then we can create a scheduler and run the inference:

graph.compile("cpu", aidge_core.dtype.float32, dims=[[1,10]])
scheduler = aidge_core.SequentialScheduler(graph)
scheduler.forward(data=[in_tensor])

for outNode in graph.get_output_nodes():
    output_aidge = np.array(outNode.get_operator().get_output(0))
    print('Aidge prediction = ', output_aidge)

Updating ONNX import

We have seen how to handle a GenericOperator in order to generate an export or run inference. However, this is not the only approach available when dealing with an operator that is not natively supported.

As mentioned earlier, the Swish function is composed of the Exp, Add, and Div operations. In this section, we will explore how to interact with the aidge_onnx library to add support for new operators. This section will also showcase the use of MetaNodes.

Creating a MetaNode

The first step is to reproduce the Swish operation using a MetaOperator.

To achieve this, we need to create a Producer Node for each constant: exp, 1, and beta. Then, we define each operation: Exp, Add, and Div. Finally, we create a GraphView, which will be embedded within a MetaOperator.

Note: The Swish computation graph begins with a branch split. To ensure a single input to the graph, we use the Identity operator.

from math import exp

def gen_swish_metaop(nb_chan, name):

    # Declaring constant values
    e_prod = aidge_core.Producer(aidge_core.Tensor(np.array([exp(1)]*nb_chan, dtype=np.float32)), "exp")
    one_prod = aidge_core.Producer(aidge_core.Tensor(np.array([1]*nb_chan, dtype=np.float32)), "one")
    beta = 0.1
    beta_prod = aidge_core.Producer(aidge_core.Tensor(np.array([-beta]*nb_chan, dtype=np.float32)), "beta")

    # Declaring operators
    mul_op = aidge_core.Mul(name=f"{name}_MUL")
    pow_op = aidge_core.Pow(name=f"{name}_POW")
    add_op = aidge_core.Add(name=f"{name}_ADD")
    div_op = aidge_core.Div(name=f"{name}_DIV")
    input_op = aidge_core.Identity(f"{name}_Input")

    # Declaring Connectors
    x = aidge_core.Connector(input_op)
    b = aidge_core.Connector(beta_prod)
    e = aidge_core.Connector(e_prod)
    o = aidge_core.Connector(one_prod)

    # Graph creation using functionnal declaration
    y = div_op(x, add_op(pow_op(e, mul_op(x, b)), o))
    swish_micro_graph = aidge_core.generate_graph([y])

    # Embedding GraphView in a MetaOperator
    swish_node = aidge_core.meta_operator(
        "Swish",
        swish_micro_graph,
        name=name
    )
    return swish_node

# Testing swich metaop
swish_node = gen_swish_metaop(10, "Test")

We can then visualize the MicroGraph of the MetaOperator Swish using Aidge Model Explorer:

import aidge_model_explorer

aidge_model_explorer.visualize(swish_node.get_operator().get_micro_graph(), "swish_micro", embed=True)

We have successfully created a function which can create a MetaOperator for the Swish function! The next step is to register this function so that it is called by the ONNX import library.

Registering new node import

Registering a new node to the ONNX import library can be easily done using the decorator function @aidge_onnx.node_import.auto_register_import.

This decorator registers the function to the dictionary of import functions aidge_onnx.node_converter.ONNX_NODE_CONVERTER_. Note that the key you should use is the ONNX name of the operator in lowercase.

NB_CHAN = 10 # TODO: Find a way to infer nb channel later ...

@aidge_onnx.node_import.auto_register_import("swish")
def import_swish(onnx_node, input_nodes, opset=None):
    node_name = onnx_node.output[0]
    return gen_swish_metaop(NB_CHAN, node_name)

Once this is done, you can use aidge_onnx.node_import.supported_operators() and check that Swish is part of the supported operators:

aidge_onnx.node_import.supported_operators()

Since Swish is now supported, it is possible to load again the ONNX file:

supported_graph = aidge_onnx.load_onnx("test_swish.onnx")

Since we have decomposed the Swish operation into atomic operators supported by Aidge, we do not need to provide a custom implementation. Instead, we can simply use the aidge_backend_cpu implementation to run inference:

data_input = aidge_core.Producer(aidge_core.Tensor(np.arange(NB_CHAN, dtype=np.float32)+1.0), "data")

data_input.add_child(supported_graph)
supported_graph.add(data_input)

data_input.get_operator().set_datatype(aidge_core.dtype.float32)

data_input.get_operator().set_backend("cpu")

supported_graph.set_datatype(aidge_core.dtype.float32)
supported_graph.set_backend("cpu")

# Create SCHEDULER
scheduler = aidge_core.SequentialScheduler(supported_graph)

# Run inference !
scheduler.forward()

for outNode in supported_graph.get_output_nodes():
    output_aidge = np.array(outNode.get_operator().get_output(0))
    print("MetaOperator output:")
    print(output_aidge)

x = np.arange(NB_CHAN, dtype=np.float32)+1.0

beta = 0.1
print("Reference output:")

print(x / (1. + np.exp(-beta*x)))