" model = helper.make_model(graph, producer_name=\"custom-model\")\n",
" model = qonnx_make_model(graph, producer_name=\"custom-model\")\n",
" return ModelWrapper(model)"
]
},
...
...
%% Cell type:markdown id: tags:
# Introduction to custom ops in FINN
Suppose that you want to introduce a new (custom) operation type into the FINN compiler. Custom operations in FINN are useful for a variety of things ranging from code generation to functional verification. This is achieved by creating a new Python module for your custom operation that fulfills certain interface specifications.
One thing to point out before we start is that **these custom operations are generic** and not really tied to e.g. Vivado HLS or few-bit quantization. As you will see in this notebook, it's possible to provide arbitrary Python/C/C++/... execution and code generation paths for custom nodes.
## The CustomOp base class
Subclasses of `CustomOp` provide a way of providing custom functionality for ONNX nodes in FINN.
This is the base class for every custom op node used in the framework, so you must create subclasses of `CustomOp` to provide execution, code generation and other functionalities in FINN.
Let's start by looking at the `CustomOp` base class itself, which lives in the `finn-base` repository. You can view it [here](https://github.com/Xilinx/finn-base/blob/dev/src/finn/custom_op/base.py). Note that the `finn` Docker container already has `finn-base` set up as a dependency.
Some points of importance:
1.`CustomOp` instances (in Python) are not meant to store any data, only provide functionality on top of data stored in ONNX. Each `CustomOp` instance has a member `self.onnx_node` which gives access to the ONNX `NodeProto` with attributes. There is also a custom attribute setter/getter system in `CustomOp` to make this process easier.
2.`CustomOp` subclasses need to implement the methods below (those not starting with underscore).
3. To be discoverable in the custom op register, `CustomOp` subclasses must set the `domain` field to the name of the Python module they appear in. For instance, to use the custom `Im2Col` op type from [here](https://github.com/Xilinx/finn-base/blob/dev/src/finn/custom_op/general/im2col.py), the ONNX node must use `domain=qonnx.custom_op.general` since its module is located at `finn/custom_op/general/im2col.py`.
%% Cell type:code id: tags:
``` python
fromqonnx.custom_op.baseimportCustomOp
dir(CustomOp)
```
%% Cell type:markdown id: tags:
## A Simple CustomOp Example
Let's make a simple CustomOp that raises its input to a given exponent (specified as attribute). For now it'll only work in Python, but later we'll add C++ execution capability too.
%% Cell type:code id: tags:
``` python
fromonnximporthelper
importnumpyasnp
classMyPythonPowerOp(CustomOp):
# here we use the CustomOp attribute system to make it easier
# to set/get custom attributes on this node
defget_nodeattr_types(self):
return{
# each entry is:
# name of attribute : (dtype, required, default value)
# dtype follows the ONNX attribute protobuf so
# "i" is int, "s" is string, "f" is float,
# "ints" is a list of integers...
# also good practice to document what each attribute does here:
# which integer power to raise the input to
"exponent":("i",True,0),
# execution mode : currently only python
"exec_mode":("s",True,"python"),
}
# return an ONNX node that has the same shape inference behavior
# here we want in shape = out shape, so we use the ONNX ReLU
# node to mimic its shape inference behavior
# we have access to the entire ModelWrapper to help make this decision
# (the parameter called model)
defmake_shape_compatible_op(self,model):
node=self.onnx_node
# make a Relu node connected to the same in-out tensors to get
# shape inference
# a general-purpose alternative is to use a Constant node that
# used for FINN DataType inference: set the output tensors' datatypes
# accordingly for this node
# here we assume input datatype = output datatype
# we have access to the entire ModelWrapper to help make this decision
# (the parameter called model)
definfer_node_datatype(self,model):
node=self.onnx_node
# data type stays the same
dtype=model.get_tensor_datatype(node.input[0])
model.set_tensor_datatype(node.output[0],dtype)
# execute this node
# context: used for both input and output, dictionary of named
# tensors
# graph: the ONNX GraphProto (ModelWrapper.graph), generally
# not needed to execute a single node
defexecute_node(self,context,graph):
exec_mode=self.get_nodeattr("exec_mode")
ifexec_mode=="python":
# get names of node input and output tensors
i_name=self.onnx_node.input[0]
o_name=self.onnx_node.output[0]
# grab input tensor from context
i_tensor=context[i_name]
# get which power to raise to from attribute
expnt=self.get_nodeattr("exponent")
# compute and put output into context
o_tensor=np.power(i_tensor,expnt)
context[o_name]=o_tensor
else:
raiseException("Only python exec_mode is supported")
# can use to do a sanity check of all the node's properties
# optional, not implemented here
defverify_node(self):
pass
```
%% Cell type:markdown id: tags:
To make sure our custom op is available, it needs to be registered. The best practice for this is to create a submodule under `finn.custom_op` which includes a `custom_op` dictionary that maps strings (op names) to classes (op implementations). Since we're in a Jupyter notebook we'll just hijack it at runtime like this:
We can see which custom ops are registered under this submodule by looking at the dictionary:
%% Cell type:code id: tags:
``` python
general.custom_op
```
%% Cell type:markdown id: tags:
## Let's Try Out our CustomOp
We'll manually build a small ONNX graph containing our node in order to try out some of the functionality. This would normally go into the unit test for this CustomOp.
The graph is built by first specifying the input/output tensor information (name, type, shape). Then,the custom node is generated; which is later used to generate the graph along the input/output tensor information. The model is built using the graph. Finally, the model is wrapped around using the ModelWrapper function from FINN.
Now, we specify the input tensor shape and we generate the graph using the function we have just created. The input tensor shape and the exponent value are passed as parameters. These parameters are used to generate our model, graph and custom node using the `MyPythonPowerOp` operation.
%% Cell type:code id: tags:
``` python
# generate a small graph with our custom op
input_shape=(1,2,4)
ret_model=make_graph(input_shape,2)
ret_model.model.graph.node
```
%% Cell type:markdown id: tags:
We generate a random tensor based on the `input_shape` defined before. See the shape and values of the `random_input` below and the datatype.
%% Cell type:code id: tags:
``` python
fromqonnx.core.datatypeimportDataType
fromqonnx.util.basicimportgen_finn_dt_tensor
# generate a random input of e.g signed 4-bit values
Just generate an Input Dictionary with the random values just generated. Then we execute the model using the model and random values just generated.
%% Cell type:code id: tags:
``` python
fromfinn.core.onnx_execimportexecute_onnx
# run with FINN's execute_onnx
inp_dict={"inp":random_input}
ret=execute_onnx(ret_model,inp_dict)
ret
```
%% Cell type:markdown id: tags:
Done! We have just executed the model that uses our custom operation. The result should be the input number to the power of 2.
%% Cell type:markdown id: tags:
## A CustomOp with C++ Generation
We can write our CustomOps in C++ for instance and generate a model the same way we have done it previously. This can be done through python bindings that let us call C++ code from python. In fact, we will compile the C++ code and execute it from python.
The following class is based on the `MyPythonPowerOp` class previously written. We are adding a new attribute `codegen_dir` into the `get_nodeattr_types` function that specifies the directory for the generated C++ code, building script and executable application.
We define a new function that `my_custom_cpp_gen` that writes the C++ code into a file and builds it. Finally the `execute_node` function is modified to support the C++ execution of the CustomOp. The `c++` branch of the if-else statements first flattens the input tensor and writes it into the "input.txt" file. Then the C++ compiled application is executed using bash commands. The application reads the ".txt" file, calculates the power value based on the exponent, and writes the result back into the "output.txt" file. Then the result of the ouput file is read and reshaped back into the original shape. Finally, the result is written into the `context` dictionary
raiseException("Only python and c++ exec_mode is supported")
# can use to do a sanity check of all the node's properties
# optional, not implemented here
defverify_node(self):
pass
```
%% Cell type:markdown id: tags:
We just register the new CustomOp the same way as we did before. Then, we create another graph using the same function `make_graph` as before. We can see the node containing the custom operation printed below.
We define a local transformation function that transforms a specific model by accessing and modifying the attributes of the specified node. It will execute the `my_custom_cpp_gen` function from the node "MyMixedPowerOp" if the "codegen_dir" is not present.
We can see that the `codegen_dir` folder contains the compile script, compiled application and the C++ source file:
%% Cell type:code id: tags:
``` python
!ls{codegen_dir}
```
%% Cell type:markdown id: tags:
Let's view the content of the C++ source file:
%% Cell type:code id: tags:
``` python
!cat{codegen_dir}/top.cpp
```
%% Cell type:markdown id: tags:
### Manually generate input and run C++ node model
We will now manually generate the input data and write it into the `input.txt` file. Then, we manually execute the compiled application and finally see the result in the `output.txt` file.
The purpose of this is mostly to show that there is no "magic" happening when FINN is executing our custom op; it's just launching a program. When debugging the execution of your custom op, it's a good idea to keep this in mind -- for instance, you can use `gdb` to debug the internals of the C++ node model here.
%% Cell type:code id: tags:
``` python
!echo"7.0 8.0 9.0">{codegen_dir}/input.txt
```
%% Cell type:code id: tags:
``` python
!cd{codegen_dir};./node_model
```
%% Cell type:code id: tags:
``` python
!cat{codegen_dir}/output.txt
```
%% Cell type:code id: tags:
``` python
!rm{codegen_dir}/*.txt
```
%% Cell type:markdown id: tags:
### Use FINN execution flow
We'll now trigger the custom node execution from inside FINN, via the custom ONNX execution capabilities which will automatically launch the appropriate handler when a custom node is encountered inside the ONNX graph, in this case launching the compiled C++ program. To do this, we will first generate a random tensor with a pre-specified tensor shape and print it.
%% Cell type:code id: tags:
``` python
# generate a random input of e.g signed 4-bit values
We set the CustomOp node attribute to execute in "Python" mode. Then, generate an input dictionay with the random input tensor and execute the transformed model using the `execute_onnx`. We print the output to see the results.
%% Cell type:code id: tags:
``` python
# run with FINN's execute_onnx, custom node will use Python execution
new_op_inst.set_nodeattr("exec_mode","python")
inp_dict={"inp":random_input}
ret=execute_onnx(mixedop_graph_new,inp_dict)
ret
```
%% Cell type:markdown id: tags:
We repeat the previous process in "c++" execution mode.
%% Cell type:code id: tags:
``` python
# run with FINN's execute_onnx, custom node will use c++ execution
This notebook should give an overview of ONNX ProtoBuf, help to create and manipulate an ONNX model and use FINN functions to work with it.
%% Cell type:markdown id: tags:
## Outline
* #### How to create a simple ONNX model
* #### How to manipulate an ONNX model
%% Cell type:markdown id: tags:
### How to create a simple ONNX model
To explain how to create an ONNX model a simple example with mathematical operations is used. All nodes are from the [standard operations library of ONNX](https://github.com/onnx/onnx/blob/master/docs/Operators.md).
First ONNX is imported, then the helper function can be used to make a node.
%% Cell type:code id: tags:
``` python
importonnx
fromqonnx.util.basicimportqonnx_make_model
Add1_node=onnx.helper.make_node(
'Add',
inputs=['in1','in2'],
outputs=['sum1'],
name='Add1'
)
```
%% Cell type:markdown id: tags:
The first attribute of the node is the operation type. In this case it is `'Add'`, so it is an adder node. Then the input names are passed to the node and at the end a name is assigned to the output.
For this example we want two other adder nodes, one abs node and the output shall be rounded so one round node is needed.
%% Cell type:code id: tags:
``` python
Add2_node=onnx.helper.make_node(
'Add',
inputs=['sum1','in3'],
outputs=['sum2'],
name='Add2',
)
Add3_node=onnx.helper.make_node(
'Add',
inputs=['abs1','abs1'],
outputs=['sum3'],
name='Add3',
)
Abs_node=onnx.helper.make_node(
'Abs',
inputs=['sum2'],
outputs=['abs1'],
name='Abs'
)
Round_node=onnx.helper.make_node(
'Round',
inputs=['sum3'],
outputs=['out1'],
name='Round',
)
```
%% Cell type:markdown id: tags:
The names of the inputs and outputs of the nodes give already an idea of the structure of the resulting graph. In order to integrate the nodes into a graph environment, the inputs and outputs of the graph have to be specified first. In ONNX all data edges are processed as tensors. So with onnx helper function tensors value infos are created for the input and output tensors of the graph. Float from ONNX is used as data type.
Now the graph can be built. First all nodes are passed. Here it is to be noted that it requires a certain sequence. The nodes must be instantiated in their dependencies to each other. This means Add2 must not be listed before Add1, because Add2 depends on the result of Add1. A name is then assigned to the graph. This is followed by the inputs and outputs.
`value_info` of the graph contains the remaining tensors within the graph. When creating the nodes we have already defined names for the inner data edges and now these are assigned tensors of the datatype float and a certain shape.
**Important**: In our example, the shape of the tensors does not change during the calculation. This is not always the case. So you have to make sure that you specify the shape correctly.
Now a model can be created from the graph and saved using the `.save` function. The model is saved in .onnx format and can be reloaded with `onnx.load()`. This also means that you can easily share your own model in .onnx format with others.
To visualize the created model, [netron](https://github.com/lutzroeder/netron) can be used. Netron is a visualizer for neural network, deep learning and machine learning models. FINN provides a utility function for visualization with netron, which we import and use in the following.
%% Cell type:code id: tags:
``` python
fromfinn.util.visualizationimportshowInNetron
```
%% Cell type:code id: tags:
``` python
showInNetron('/tmp/simple_model.onnx')
```
%% Cell type:markdown id: tags:
Netron also allows you to interactively explore the model. If you click on a node, the node attributes will be displayed.
In order to test the resulting model, a function is first written in Python that calculates the expected output. Because numpy arrays are to be used, numpy is imported first.
%% Cell type:code id: tags:
``` python
importnumpyasnp
defexpected_output(in1,in2,in3):
sum1=np.add(in1,in2)
sum2=np.add(sum1,in3)
abs1=np.absolute(sum2)
sum3=np.add(abs1,abs1)
returnnp.round(sum3)
```
%% Cell type:markdown id: tags:
Then the values for the three inputs are calculated. Random numbers are used.
We can easily pass the values to the function we just wrote to calculate the expected result. For the created model the inputs must be summarized in a dictionary, which is then passed on to the model.
%% Cell type:code id: tags:
``` python
input_dict={}
input_dict["in1"]=in1_values
input_dict["in2"]=in2_values
input_dict["in3"]=in3_values
```
%% Cell type:markdown id: tags:
To run the model and calculate the output, [onnxruntime](https://github.com/microsoft/onnxruntime) can be used. ONNX Runtime is a performance-focused complete scoring engine for ONNX models from Microsoft. The `.InferenceSession` function is used to create a session of the model and `.run` is used to execute the model.
The input values are also transferred to the reference function. Now the output of the execution of the model can be compared with that of the reference.
print("The output of the ONNX model is: \n{}".format(output[0]))
print("\nThe output of the reference function is: \n{}".format(ref_output))
if (output[0]==ref_output).all():
print("\nThe results are the same!")
else:
raiseException("Something went wrong, the output of the model doesn't match the expected output!")
```
%% Cell type:markdown id: tags:
Now that we have verified that the model works as we expected it to, we can continue working with the graph.
%% Cell type:markdown id: tags:
### How to manipulate an ONNX model
In the model there are two successive adder nodes. An adder node in ONNX can only add two inputs, but there is also the [**sum**](https://github.com/onnx/onnx/blob/master/docs/Operators.md#Sum) node, which can process more than two inputs. So it would be a reasonable change of the graph to combine the two successive adder nodes to one sum node.
%% Cell type:markdown id: tags:
In the following we assume that we do not know the appearance of the model, so we first try to identify whether there are two consecutive adders in the graph and then convert them into a sum node.
Here we make use of FINN. FINN provides a thin wrapper around the model which provides several additional helper functions to manipulate the graph. The so called `ModelWrapper` can be found in the QONNX repository which contains a lot of functionality that is used by FINN, you can find it [here](https://github.com/fastmachinelearning/qonnx/blob/main/src/qonnx/core/modelwrapper.py).
%% Cell type:code id: tags:
``` python
fromqonnx.core.modelwrapperimportModelWrapper
finn_model=ModelWrapper(onnx_model)
```
%% Cell type:markdown id: tags:
As explained in the previous section, it is important that the nodes are listed in the correct order. If a new node has to be inserted or an old node has to be replaced, it is important to do that in the appropriate position. The following function serves this purpose. It returns a dictionary, which contains the node name as key and the respective node index as value.
%% Cell type:code id: tags:
``` python
defget_node_id(model):
node_index={}
node_ind=0
fornodeinmodel.graph.node:
node_index[node.name]=node_ind
node_ind+=1
returnnode_index
```
%% Cell type:markdown id: tags:
The function scans the list of nodes and stores a run index (`node_ind`) as node index in the dictionary for every node name.
Another helper function is being implemented that searches for adder nodes in the graph and returns the found nodes. This is needed to determine if and which adder nodes are in the given model.
%% Cell type:code id: tags:
``` python
defidentify_adder_nodes(model):
add_nodes=[]
fornodeinmodel.graph.node:
ifnode.op_type=="Add":
add_nodes.append(node)
returnadd_nodes
```
%% Cell type:markdown id: tags:
The function iterates over all nodes of the model and if the operation type is `"Add"` the node will be stored in `add_nodes`. At the end `add_nodes` is returned.
If we apply this to our model, three nodes should be returned.
%% Cell type:code id: tags:
``` python
add_nodes=identify_adder_nodes(finn_model)
fornodeinadd_nodes:
print("Found adder node: {}".format(node.name))
```
%% Cell type:markdown id: tags:
Among other helper functions, `ModelWrapper` offers two functions that can help to determine the preceding and succeeding node of a node: `find_direct_successors` and `find_direct_predecessors`. So we can use one of them to define a function to find adder pairs.
%% Cell type:code id: tags:
``` python
defadder_pair(model,node):
adder_pairs=[]
node_pair=[]
successor_list=model.find_direct_successors(node)
forsuccessorinsuccessor_list:
ifsuccessor.op_type=="Add":
node_pair.append(node)
node_pair.append(successor)
adder_pairs.append((node_pair))
node_pair=[]
returnadder_pairs
```
%% Cell type:markdown id: tags:
The function gets a node and the model as input. Two empty lists are created to be filled with a list of adder node pairs that can be returned as result of the function. Then the function `find_direct_successors` is used to return all of the successors of the node. If one of the successors is an adder node, the node is saved in `node_pair` together with the successive adder node and put in the list `adder_pairs`. Then the temporary list is cleaned and can be filled with the next adder node pair. Since it is theoretically possible for an adder node to have more than one subsequent adder node, a list of lists is created. This list of the node with all its successive adder nodes is returned.
So now we can find out which adder node has an adder node as successor. Since the model is known, one adder pair (Add1+Add2) should be found when applying the function to the previously determined adder node list (`add_nodes`).
%% Cell type:code id: tags:
``` python
fornodeinadd_nodes:
add_pairs=adder_pair(finn_model,node)
iflen(add_pairs)!=0:
foriinrange(len(add_pairs)):
substitute_pair=add_pairs[i]
print("Found following pair that could be replaced by a sum node:")
fornode_pairinadd_pairs:
fornodeinnode_pair:
print(node.name)
```
%% Cell type:markdown id: tags:
Now that the pair to be replaced has been identified (`substitute_pair`), a sum node can be instantiated and inserted into the graph at the correct position.
First of all, the inputs must be determined. For this the adder nodes inputs are used minus the input, which corresponds to the output of the other adder node.
print("The new node gets the following inputs: \n{}".format(input_list))
```
%% Cell type:markdown id: tags:
The output of the sum node matches the output of the second adder node and can therefore be taken over directly.
%% Cell type:code id: tags:
``` python
sum_output=substitute_pair[1].output[0]
```
%% Cell type:markdown id: tags:
The sum node can be created with this information.
%% Cell type:code id: tags:
``` python
Sum_node=onnx.helper.make_node(
'Sum',
inputs=input_list,
outputs=[sum_output],
name="Sum"
)
```
%% Cell type:markdown id: tags:
The node can now be inserted into the graph and the old nodes are removed.
%% Cell type:code id: tags:
``` python
node_ids=get_node_id(finn_model)
node_ind=node_ids[substitute_pair[0].name]
graph.node.insert(node_ind,Sum_node)
fornodeinsubstitute_pair:
graph.node.remove(node)
```
%% Cell type:markdown id: tags:
To insert the node in the right place, the index of the first node of the substitute_pair is used as node index for the sum node and embedded into the graph using `.insert`. Then the two elements in `substitute_pair` are deleted using `.remove`. `.insert` and `.remove` are functions provided by ONNX.
%% Cell type:markdown id: tags:
The new graph is saved as ONNX model and can be visualized with Netron.
Through the visualization it can already be seen that the insertion was successful, but it is still to be checked whether the result remains the same. Therefore the result of the reference function written in the previous section is used and the new model with the input values is simulated. At this point onnxruntime can be used again. The simulation is analogous to the one of the first model in the previous section.
