An add
OP is used as an example in this
section. For more Python examples, refer to TensorFlow2 and PyTorch examples.
- Create a Python package and a module
file.
% mkdir -p /tmp/demo_add_op/ % cd /tmp/demo_add_op/ % mkdir vart_op_imp/ % touch vart_op_imp/__init__.py % touch vart_op_imp/add.py
Note:- The folder name should be vart_op_imp.
- If you have more than one custom OP, create them separately and place them in the vart_op_imp folder.
- Update
add.py
as shown in the following code snippet.import numpy as np class add: def __init__(self, op): pass def calculate(self, output, input): np_output = np.array(output, copy=False) L = len(input) if L == 0: return np_input_0 = np.array(input[0], copy=False) np.copyto(np_output, np_input_0) for i in range(1, L): np_output = np.add(np_output, np.array(input[i], copy=False), out=np_output)
You can also use the following simplified version.
import numpy as np class add: def __init__(self, op): pass def calculate(self, output, input): np_output = np.array(output, copy=False) L = len(input) assert L == 2 np_input_0 = np.array(input[0], copy=False) np_input_1 = np.array(input[1], copy=False) np_output = np.add(np_input_0, np_input_1, out=np_output)
- Install the op as shown in the following code
snippet.
% mkdir -p lib % ln -sf ~/.local/Ubuntu.18.04.x86_64.Debug/lib/libvart_op_imp_python-cpu-op.so lib/libvart_op_imp_add.so % ls -l lib % mkdir -p dump % env LD_LIBRARY_PATH=$HOME/.local/Ubuntu.18.04.x86_64.Debug/lib:$PWD/lib $HOME/.local/Ubuntu.18.04.x86_64.Debug/share/vitis_ai_library/test/cpu_task/test_op_imp --graph /tmp/add.xmodel --op "add_op"
For Edge, you can execute the following command to install and test the op.
# ls -sf /usr/lib/libvart_op_imp_python-cpu-op.so /usr/lib/libvart_op_imp_add.so
# cp -r vart_op_imp /usr/lib.python3.9/site-packages
# xdputil run_op add.xmodel add_op -r ref -d dump
Similar to the C++ interface, the Python module must have a class whose name
is the same as the type of xir::Op. In this example, add
is used. The class must have a constructor with a single argument
op in addition to self. It is an XIR::Op. Refer to the XIR Python API for more
information. Similarly, the class add
must have a member function
called calculate
, in addition to the self
argument. The first argument must be named output
, and following argument names must comply with
the XIR::OpDef associated with the XIR::Op, refer to XIR API for more detail.