Add a custom operator to the C++ export

The main objective of this tutorial is to demonstrate the toolchain to detect unsupported operators and add them in an export module. For this tutorial, we use the CPP export module aidge_export_cpp to demonstrate the toolchain.

# First import some utility methods used in the tutorial:
import sys, os
sys.path.append(os.path.abspath(os.path.join('../../..')))
import tuto_utils

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[1], line 4
      2 import sys, os
      3 sys.path.append(os.path.abspath(os.path.join('../../..')))
----> 4 import tuto_utils

ModuleNotFoundError: No module named 'tuto_utils'

Import Aidge

import aidge_core
import aidge_backend_cpu
import aidge_export_cpp
import aidge_onnx
import numpy as np
import os
import requests

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[2], line 3
      1 import aidge_core
      2 import aidge_backend_cpu
----> 3 import aidge_export_cpp
      4 import aidge_onnx
      5 import numpy as np

ModuleNotFoundError: No module named 'aidge_export_cpp'

Load ONNX model

file_url = "https://huggingface.co/EclipseAidge/LeNet/resolve/main/lenet_mnist.onnx?download=true"
file_path = "lenet_mnist.onnx"
aidge_core.utils.download_file(file_path, file_url)

model = aidge_onnx.load_onnx("lenet_mnist.onnx")

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[4], line 1
----> 1 model = aidge_onnx.load_onnx("lenet_mnist.onnx")

NameError: name 'aidge_onnx' is not defined

# Remove Flatten node, useless in the CPP export
aidge_core.remove_flatten(model)

digit = np.load("digit.npy")

# Create Producer Node for the Graph
# Note: This means the graph will have no input!
input_node = aidge_core.Producer(aidge_core.Tensor(digit), "input")
input_node.add_child(model)
model.add(input_node)

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[5], line 2
      1 # Remove Flatten node, useless in the CPP export
----> 2 aidge_core.remove_flatten(model)
      4 digit = np.load("digit.npy")
      6 # Create Producer Node for the Graph
      7 # Note: This means the graph will have no input!

NameError: name 'model' is not defined

Replace ReLU operators by Swish operators

Let’s say you want to replace ReLU with another activation like Switch.

# Forward the dimensions in the graph in order to get the size for the beta vector of the Swish
model.forward_dims()

# Use GraphMatching to replace ReLU with Swish
matches = aidge_core.SinglePassGraphMatching(model).match("ReLU")
print('Number of match : ', len(matches))

switch_id = 0
for i, match in enumerate(matches):
    print('Match ', i)
    node_ReLU = match.graph.root_node()

    print('Replacing ', node_ReLU.type(), ' : ', node_ReLU.name())

    # We instanciate Swish as a generic operator
    node_swish = aidge_core.GenericOperator("Swish", nb_data=1, nb_param=0, nb_out=1, name=f"swish_{switch_id}")
    node_swish.get_operator().attr.betas = [1.0]*node_ReLU.get_operator().get_input(0).dims()[1]

    print('Replacing ', node_ReLU.type(), ' : ', node_ReLU.name(), ' with ' , node_swish.name())
    aidge_core.GraphView.replace(set([node_ReLU]), set([node_swish]))

    switch_id+=1

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[6], line 2
      1 # Forward the dimensions in the graph in order to get the size for the beta vector of the Swish
----> 2 model.forward_dims()
      4 # Use GraphMatching to replace ReLU with Swish
      5 matches = aidge_core.SinglePassGraphMatching(model).match("ReLU")

NameError: name 'model' is not defined

model.save("myModel")
tuto_utils.visualize_mmd("myModel.mmd")

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[7], line 1
----> 1 model.save("myModel")
      2 tuto_utils.visualize_mmd("myModel.mmd")

NameError: name 'model' is not defined

Schedule the graph

Add the function to specify how Swish activation transforms the dimensions. This forward_dims function is required to perform a sequential scheduling.

class GenericImpl(aidge_core.OperatorImpl): # Inherit OperatorImpl to interface with Aidge !
    def __init__(self, op: aidge_core.Operator):
        aidge_core.OperatorImpl.__init__(self, op, 'cpu')
    # no need to define forward() function in python as we do not intend to run a scheduler on the model

for node in model.get_nodes():
    if node.type() == "Swish":
        node.get_operator().set_forward_dims(lambda x: x) # to propagate dimensions in the model

for node in model.get_nodes():
    node.get_operator().set_impl(GenericImpl(node.get_operator())) # Setting implementation

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[8], line 6
      3         aidge_core.OperatorImpl.__init__(self, op, 'cpu')
      4     # no need to define forward() function in python as we do not intend to run a scheduler on the model
----> 6 for node in model.get_nodes():
      7     if node.type() == "Swish":
      8         node.get_operator().set_forward_dims(lambda x: x) # to propagate dimensions in the model

NameError: name 'model' is not defined

scheduler = aidge_core.SequentialScheduler(model)
model.compile("cpu", aidge_core.dtype.float32, dims=[[1, 1, 28, 28]])
scheduler.generate_scheduling()
s = scheduler.get_sequential_static_scheduling()

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[9], line 1
----> 1 scheduler = aidge_core.SequentialScheduler(model)
      2 model.compile("cpu", aidge_core.dtype.float32, dims=[[1, 1, 28, 28]])
      3 scheduler.generate_scheduling()

NameError: name 'model' is not defined

Add Swish to the CPP export support

# Note: we register a GenericOperator so we need to use ``register_generic``.
# For registering an existing operator use ``register``
# For registering a MetaOperator use ``register_metaop``
@aidge_export_cpp.ExportLibCpp.register_generic("Swish", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.float32)))
class SwishCPP(aidge_core.export_utils.ExportNodeCpp):
    def __init__(self, node, mem_info):
        super().__init__(node, mem_info)
        self.config_template = "swish_export_files/swish_config.jinja"

        self.forward_template = "swish_export_files/swish_forward.jinja"
        self.include_list = []
        self.kernels_to_copy = [
            "swish_export_files/swish_kernel.hpp",
        ]

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[10], line 4
      1 # Note: we register a GenericOperator so we need to use ``register_generic``.
      2 # For registering an existing operator use ``register``
      3 # For registering a MetaOperator use ``register_metaop``
----> 4 @aidge_export_cpp.ExportLibCpp.register_generic("Swish", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.float32)))
      5 class SwishCPP(aidge_core.export_utils.ExportNodeCpp):
      6     def __init__(self, node, mem_info):
      7         super().__init__(node, mem_info)

NameError: name 'aidge_export_cpp' is not defined

aidge_core.export_utils.scheduler_export(
        scheduler,
        "myexport",
        aidge_export_cpp.ExportLibCpp,
        memory_manager=aidge_core.mem_info.generate_optimized_memory_info,
        memory_manager_args={"stats_folder": "myexport/stats", "wrapping": False }
)

aidge_core.export_utils.generate_main_cpp("myexport", model)

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[11], line 2
      1 aidge_core.export_utils.scheduler_export(
----> 2         scheduler,
      3         "myexport",
      4         aidge_export_cpp.ExportLibCpp,
      5         memory_manager=aidge_core.mem_info.generate_optimized_memory_info,
      6         memory_manager_args={"stats_folder": "myexport/stats", "wrapping": False }
      7 )
      9 aidge_core.export_utils.generate_main_cpp("myexport", model)

NameError: name 'scheduler' is not defined

!tree myexport

/usr/bin/sh: 1: tree: not found

!cd myexport && make

/usr/bin/sh: 1: cd: can't cd to myexport

!./myexport/bin/run_export

/usr/bin/sh: 1: ./myexport/bin/run_export: not found