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.

For this tutorial, we will need the following libraries:

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

Setting up the notebook

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 this piece of code :

import os
import requests

def download_material(path: str) -> None:
    if not os.path.isfile(path):
        response = requests.get("https://gitlab.eclipse.org/eclipse/aidge/aidge/-/raw/dev/examples/tutorials/Onnx_tutorial/"+path+"?ref_type=heads")
        if response.status_code == 200:
            with open(path, 'wb') as f:
                f.write(response.content)
            print("File downloaded successfully.")
        else:
            print("Failed to download file. Status code:", response.status_code)

# Download onnx model file
download_material("test_swish.onnx")

File downloaded successfully.

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")

- Swish0 (Swish | GenericOperator)
        - beta : [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]

The swish operator is not supported and thus is loaded as a GenericOperator. This mechanism allow to load 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)

The graph is not fully supported by aidge !

Native operators: 0 (0 types)
Generic operators: 1 (1 types)
- Swish: 1
Native types coverage: 0.0% (0/1)
Native operators coverage: 0.0% (0/1)

(defaultdict(int, {}), defaultdict(int, {'Swish': 1}))

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

Working with generic operator

Indeed, using the python library, we can work with GenericOperator !

For this we will begin with retrieving the operator:

swish_node = graph.get_nodes().pop()   # get_nodes return 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 in/out dimensions.

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

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

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

swish_op.set_forward_dims(lambda x: x)

Providing an implementation

If we want to run an inference, we need to provide an implementation. Which means define the forward function.

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

So 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}")

Random input:
[[ 0.37969178  1.4566833  -1.249093    0.27077118 -1.1967081   1.0207533
   1.0625035   0.5436063  -0.38500533 -0.36396742]]

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)

Aidge prediction =  [[ 0.22546051  1.18140636 -0.27836935  0.15360367 -0.27770969  0.75037543
   0.78961928  0.34391302 -0.15589645 -0.14922646]]

Updating ONNX import

We have seen how to handle GenericOperator in order to generate an export or run inference. However, this is not the only approach we can have to support an unsupported operator.

As stated above, the Swish function is the composition of an Exp, Add and Div. In this section, we will see how we can interact with the aidge_onnx library in order to add the 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.

For this we will need to create a Producer Node for each constant: exp, 1 and beta. Then define each function Exp, Add and Div. And then create a GraphView that we will embedded in a MetaOperator.

Note: The swish computation graph begin with a branch split, so to have one input I use the operator Identity.

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(2, 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])

    # Saving micrograph for visualisation
    swish_micro_graph.save("swish_micro")

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

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

We can then visualize the micro graph of the macro operator swish using mermaid:

import base64
from IPython.display import Image, display
import matplotlib.pyplot as plt

def visualize_mmd(path_to_mmd):
  with open(path_to_mmd, "r") as file_mmd:
    graph_mmd = file_mmd.read()

  graphbytes = graph_mmd.encode("ascii")
  base64_bytes = base64.b64encode(graphbytes)
  base64_string = base64_bytes.decode("ascii")
  display(Image(url=f"https://mermaid.ink/img/{base64_string}"))


visualize_mmd("swish_micro.mmd")

We have successfully created a function which can create a MetaOperator for the Swish function !

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 will register the function to the dictionary of import function 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()

['softmax',
 'div',
 'reducemean',
 'pow',
 'mul',
 'tanh',
 'sigmoid',
 'leakyrelu',
 'slice',
 'sqrt',
 'globalaveragepool',
 'concat',
 'constant',
 'averagepool',
 'split',
 'relu',
 'gemm',
 'gather',
 'identity',
 'matmul',
 'batchnorm',
 'batchnormalization',
 'transpose',
 'maxpool',
 'lstm',
 'erf',
 'add',
 'sub',
 'shape',
 'conv',
 'reshape',
 'dropout',
 'swish']

Since swish is supported we can load again the onnx:

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

Since we have decomposed the Swish operation in atomic operator supported by Aidge, we don’t need to provide an implementation and instead we can just use the aidge_backend_cpu implementation to run an 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)))

MetaOperator output:
[0.5249792 1.0996679 1.7233275 2.3947506 3.1122966 3.8739378 4.6773143
 5.519796  6.3985453 7.310586 ]
Reference output:
[0.5249792 1.099668  1.7233275 2.3947506 3.1122966 3.8739378 4.6773148
 5.519796  6.3985453 7.310586 ]