Write OP Implementation in Python - 2.0 English

Vitis AI Library User Guide (UG1354)

Document ID
UG1354
Release Date
2022-01-20
Version
2.0 English

An add OP is used as an example in this section.

  1. 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
  2. Update add.py as below.
    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)
  3. Install the op as below.
    % 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.8/site-packages
# xdputil run_op add.xmodel add_op -r ref -d dump
Note: Before you execute the above commands, copy add.xmodel, ref folder with sample input files and vart_op_imp folder to the board.

Similar to the C++ interface, the Python module must have a class whose name is the same as the xir::Op's type. It is add in the example. 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 detail. Similarly, the class add must have a member function calculate, in addition to 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.

Note: A symbolic link to libvart_op_imp_python-cpu-op.so named libvart_op_imp_add.so is created. libvart_op_imp_python-cpu-op.so is a bridge between Python and C++. From the C++ side, it searches for libvart_op_imp_add.so and finds the libvart_op_imp_python-cpu-op.so. In libvart_op_imp_python_cpu_op.so, the Python module name vart_op_imp.add is imported and Python searches for the module as usual.