Making an onnx model with custom/contrib OP (like QAttention) using onnx python hooks?

[sunnyshu-intel]

Hi, am new to onnx, I want to generate a model/graph which has only one QAttention node/layer, to test QAttention with different outputs and inputs. I have a script which helps me to generate models/graphs which consist of only node/layer which I want to test, for example, MatMul + Add as a node, coded as follows.

import onnx
from onnx import helper, shape_inference
from onnx import numpy_helper
from onnx import AttributeProto, TensorProto, GraphProto
import os
import numpy.matlib
import numpy as np

_this_dir = os.path.abspath(os.path.dirname(__file__))

def _onnx_export(graph_def, relative_path, verbose=False):
  model = helper.make_model(graph_def, producer_name='makalini', opset_imports=[helper.make_operatorsetid("", 11)])
  onnx.checker.check_model(model)
  inferred_model = shape_inference.infer_shapes(model)
  onnx.checker.check_model(inferred_model)
  model_path = os.path.join(_this_dir, relative_path)
  os.makedirs(os.path.dirname(model_path), exist_ok=True)
  onnx.save_model(model, model_path)
  if verbose:
    print()
    print(inferred_model)
    import onnxruntime as rt
    rt.InferenceSession(model_path)

def matmul_add():
    model_name = 'matmul_add'
    x = np.array([1,2,3,4,5,6]).reshape([2,3]).astype(np.float32)
    w = np.array([1,2,3,-1,-2,-3]).reshape([3,2]).astype(np.float32)
    a = np.array([0.2, 0.5]).reshape([2]).astype(np.float32)

    matmul1 = np.matmul(x, w)
    y = np.add(matmul1, a)

    graph_def = helper.make_graph(
        nodes = [
          helper.make_node(op_type = "MatMul", inputs = ['x', 'w'], outputs = ['matmul1'], name = 'matmul_1'),
          helper.make_node(op_type = "Add", inputs = ['matmul1', 'a'], outputs = ['y'], name = 'add_1')
        ],
        name = model_name,
        inputs = [
          helper.make_tensor_value_info('x', TensorProto.FLOAT, np.shape(x)),
          helper.make_tensor_value_info('w', TensorProto.FLOAT, np.shape(w)),
          helper.make_tensor_value_info('a', TensorProto.FLOAT, np.shape(a)),
        ],
        outputs = [
          helper.make_tensor_value_info('y', TensorProto.FLOAT, np.shape(y)),
        ],
        initializer = [
        ]
    )
    _onnx_export(graph_def, os.path.join(model_name, model_name + '.onnx'))

def generate_all():
  matmul_add()

if __name__ == '__main__':
  generate_all()

I want to do same with the QAttention (or similar ops). Below is the code I trying to generate a model consist of only QAttention node, but it always gives me error. I am not sure what am doing wrong or what I should be doing extra to do the same, Please help.

NOTE: In the below code just suppose we are giving the the correct shape input, we can share code regarding the same if needed.

generate_qattention.py

import numpy as np
from enum import Enum
import onnx
from onnx import helper, shape_inference
from onnx import numpy_helper
from onnx import AttributeProto, TensorProto, GraphProto
import os
import numpy as np

def onnx_export(graph_def, relative_path, verbose=False):

  # A Function to check if model is valid or not
  def is_valid(onnx_model):
    try:
        onnx.checker.check_model(onnx_model)
    except onnx.checker.ValidationError as e:
        print('The model is invalid: %s' % e)
        return False
    else:
        print('The model is valid!')
        return True
  
  # Make model, check validity and save if model is valid
  model = helper.make_model(graph_def, producer_name='makalini', opset_imports=[helper.make_operatorsetid("", 11)])
  if is_valid(model):
    model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), relative_path)
    os.makedirs(os.path.dirname(model_path), exist_ok=True)
    onnx.save_model(model, model_path)
    
    # Print inferred model
    if verbose:
      inferred_model = shape_inference.infer_shapes(model)
      if is_valid(inferred_model):
        print()
        print(inferred_model)
        import onnxruntime as rt
        rt.InferenceSession(model_path)


class DATA(Enum):
    INPUT = 0
    WEIGHTS = 1
    BIAS = 2
    INPUT_SCALE = 3
    WEIGHTS_SCALE = 4
    MASK_INDEX = 5
    INPUT_ZP = 6
    WEIGHTS_ZP = 7
    OUTPUT = 8

def qattention(x, w, b, x_scale, w_scale, m, x_zero_point, w_zero_point, y):

    # Set Prerequisite
    model_name = 'qattention'

    # Make Graph
    G = helper.make_graph(
        nodes = [
          helper.make_node(op_type = "QAttention",
                           inputs = [DATA.INPUT.name,
                                     DATA.WEIGHTS.name,
                                     DATA.BIAS.name,
                                     DATA.INPUT_SCALE.name,
                                     DATA.WEIGHTS_SCALE.name,
                                     DATA.MASK_INDEX.name,
                                     DATA.INPUT_ZP.name,
                                     DATA.WEIGHTS_ZP.name],
                           outputs = [DATA.OUTPUT.name],
                           name = 'qattention_1'),
        ],
        name = model_name,
        inputs = [
          helper.make_tensor_value_info(DATA.INPUT.name, TensorProto.FLOAT, x.shape),
          helper.make_tensor_value_info(DATA.WEIGHTS.name, TensorProto.FLOAT, w.shape),
          helper.make_tensor_value_info(DATA.BIAS.name, TensorProto.FLOAT, b.shape),
          helper.make_tensor_value_info(DATA.INPUT_SCALE.name, TensorProto.FLOAT, x_scale.shape),
          helper.make_tensor_value_info(DATA.WEIGHTS_SCALE.name, TensorProto.FLOAT, w_scale.shape),
          helper.make_tensor_value_info(DATA.MASK_INDEX.name, TensorProto.INT64, m.shape),
          helper.make_tensor_value_info(DATA.INPUT_ZP.name, TensorProto.INT32, x_zero_point.shape),
          helper.make_tensor_value_info(DATA.WEIGHTS_ZP.name, TensorProto.INT32, w_zero_point.shape),
        ] ,
        outputs = [
          helper.make_tensor_value_info(DATA.OUTPUT.name, TensorProto.FLOAT, y.shape),
        ],
        initializer = [
        ]
    )
    onnx_export(G, os.path.join(model_name, model_name + '.onnx'))

    return

qattention()

Terminal Output

PS C:\Users\sunnyshu> python -u "c:\ss\workspace\generate_model\generate_qattention.py"
The model is invalid: No Op registered for QAttention with domain_version of 11

==> Context: Bad node spec for node. Name: qattention_1 OpType: QAttention