netrans/bin/custom/custom_sample_1.py

import numpy as np
from netranslib.custom_interface import *

# Layout perm forward callback function
def forward_callback_cond(net, lid):
    # This is a generic method to check condition when forward optimize perm.
    ret = forward_perm_check(net, lid)
    print("forward_callback_cond")
    return ret

def forward_callback_move(net, lid, fmt):
    # This is a generic method to forward optimize perm.
    forward_perm_move(net, lid, fmt)
    print("forward_callback_move")

def forward_callback_cvt(net, lid, fmt):
    # Convert parameters if you need when forward optimize perm.
    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("forward_callback_cvt")

def forward_callback_data(net, lid):
    # Transpose data if you need when forward optimize perm.
    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("forward_callback_data")

# Layout perm backward callback function
def backward_callback_cond(net, lid):
    # This is a generic method to check condition when backward optimize perm.
    ret = backward_perm_check(net, lid)
    print("backward_callback_cond")
    return ret

def backward_callback_move(net, lid, fmt):
    # This is a generic method to backward move perm.
    backward_perm_move(net, lid, fmt)
    print("backward_callback_move")

def backward_callback_cvt(net, lid, fmt):
    # Convert parameters if you need when backward optimize perm.
    l = net.get_layer(lid)
    axis = l.get_params()['axis']
    perm = get_output_perm(l.get_output(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("backward_callback_cvt")

def backward_callback_data(net, lid):
    # Transpose data if you need when backward optimize perm.
    l = net.get_layer(lid)
    sample_data = l.get_const_var('sample_data')
    perm = get_output_perm(l.get_output(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("backward_callback_data")

# Quantization callback function
def quantize_check_callback(l):
    # Add your check condition if you need.
    print("quantize_check_callback")
    return True

class CustomSample1(CustomOp):

    op = 'custom_sample_1'

    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):
        out = input_tensor[0]
        return [out]

    def setup_same_min_max(self):
        cfg = QuantMinMaxConfig()
        cfg.same_min_max = ['in0', 'out0']
        cfg.cond = quantize_check_callback
        return cfg

    def layout_perm_optimize_forward(self):
        cfg = MarkOpConfig()
        cfg.cond = forward_callback_cond
        cfg.move = forward_callback_move
        cfg.cvt = forward_callback_cvt
        cfg.data = forward_callback_data
        return cfg

    def layout_perm_optimize_backward(self):
        cfg = MarkOpConfig()
        cfg.cond = backward_callback_cond
        cfg.move = backward_callback_move
        cfg.cvt = backward_callback_cvt
        cfg.data = backward_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)