netrans/bin/custom/custom_sample_0.py

import numpy as np
from netranslib.custom_interface import *

# Layout convert callback function nhwc to nchw
def t2c_callback_cond(net, lid):
    # Add your check condition if you need when layout from nhwc to nchw.
    print("t2c_callback_cond")
    return True

def t2c_callback_proc(net, lid):
    # Convert all inputs layout nhwc to nchw.
    convert_t2c_inputs(net, lid, perm=[0,3,1,2])

    # Convert all outputs layout nchw to nhwc.
    convert_t2c_outputs(net, lid, perm=[0,2,3,1])
    print("t2c_callback_proc")

def t2c_callback_cvt(net, lid):
    # Convert parameters if you need when layout from nhwc to nchw.
    l = net.get_layer(lid)
    axis = l.get_params()['axis']
    perm = get_input_perm(l.get_input(0))

    # Please fill your function to convert the parameters if you need.
    # The following is just an example
    new_axis = perm[perm[axis]]

    l.put_param('axis', new_axis)
    print("t2c_callback_cvt")

def t2c_callback_data(net, lid):
    # Transpose data if you need when layout from nhwc to nchw.
    l = net.get_layer(lid)
    sample_data = l.get_const_var('sample_data')
    perm = get_input_perm(l.get_input(0))

    # Please fill your function to transpose the data if you need.
    # The 'sample_data' is a ndarray.
    new_data = sample_data

    l.put_const_tensor('sample_data', new_data)
    print("t2c_callback_data")

# Layout convert callback function nchw to nhwc
def c2t_callback_cond(net, lid):
    # Add your check condition to do layout convert.
    print("c2t_callback_cond")
    return True

def c2t_callback_proc(net, lid):
    # Convert all inputs layout nchw to nhwc.
    convert_t2c_inputs(net, lid, perm=[0,2,3,1])

    # Convert all outputs layout nhwc to nchw.
    convert_t2c_outputs(net, lid, perm=[0,3,1,2])
    print("c2t_callback_proc")

def c2t_callback_cvt(net, lid):
    # Convert parameters if you need when layout from nchw to nhwc.
    l = net.get_layer(lid)
    axis = l.get_params()['axis']
    perm = get_input_perm(l.get_input(0))

    # Please fill your function to convert the parameters if you need.
    # The following is just an example.
    new_axis = perm[perm[axis]]

    l.put_param('axis', new_axis)
    print("c2t_callback_cvt")

def c2t_callback_data(net, lid):
    # Transpose data if you need when layout from nhwc to nchw.
    l = net.get_layer(lid)
    sample_data = l.get_const_var('sample_data')
    perm = get_input_perm(l.get_input(0))

    # Please fill your function to transpose the data if you need.
    # The 'sample_data' is a ndarray.
    new_data = sample_data

    l.put_const_tensor('sample_data', new_data)
    print("c2t_callback_data")

class CustomSample0(CustomOp):

    op = 'custom_sample_0'

    def_input  = [IoMap('in0', 'in', 'input port')]
    def_output = [IoMap('out0', 'out', 'output port')]

    coef = ['sample_data']

    def_param = [
        DefParam('axis', -1, False),
    ]

    def get_variable_shape(self, coef):
        if coef != 'sample_data':
            al.e('Unsupported coef:{}'.format(coef))
        shape = [1]
        return shape

    def compute_shape_nchw(self):
        return self.compute_shape_nhwc()

    def compute_shape_nhwc(self):
        out_shape = self.get_input(0).shape.copy()
        return [out_shape]

    def compute_out_tensor(self, tensor, input_tensor):
        one = np.ones([1]).astype(np.float32)
        out = input_tensor[0] + one
        return [out]

    def setup_qnt_tensors(self):
        cfg = DefaultQuantTensors()
        cfg.tensors = ['out0']
        cfg.rules = ['quantize_inputs']
        return cfg

    def layout_nchw_to_nhwc(self):
        cfg = MarkOpConfig()
        cfg.proc = c2t_callback_proc
        cfg.cond = c2t_callback_cond
        cfg.cvt = c2t_callback_cvt
        cfg.data = c2t_callback_data
        return cfg

    def layout_nhwc_to_nchw(self):
        cfg = MarkOpConfig()
        cfg.proc = t2c_callback_proc
        cfg.cond = t2c_callback_cond
        cfg.cvt = t2c_callback_cvt
        cfg.data = t2c_callback_data
        return cfg

    # Please fill your custom optimize when export case.
    def export_optimize(self, net):
        # process parameters
        # Such as reverse 'axis', because driver list is whcn, netrans is nchw
        axis = self.get_params()['axis']
        dim_num = self.get_input(0).shape.rank - 1
        new_axis = dim_num - axis
        self.put_param('axis', new_axis)

        # process data
        sample_data = self.get_const_var('sample_data')
        new_data = sample_data
        self.put_const_tensor('sample_data', new_data)

    # Please override this function if you want to export a ovxlib builtin custom op
    def export_node_map(self, l):
        NODE_TYPE = 'VSI_NN_OP_CUSTOM_SAMPLE'
        p = l.get('parameters', None)
        const_tensor = list()
        body = list()

        fill_param(body, "custom_sample.axis", int(p['axis']))
        const_tensor.append((p['data_sample_data'], 1, 'data_sample_data'))
        return NODE_TYPE, body, const_tensor