Register Custom Operation - 2.0 English

Vitis AI User Guide (UG1414)

Document ID
Release Date
2.0 English
In order to convert a quantized model to an xmodel,vai_q_pytorch provides a decorator to register an operation or a group of operations as a custom operation which is unknown for XIR.
# Decorator API
def register_custom_op(op_type: str, attrs_list: Optional[List[str]] = None):
  """The decorator is used to register the function as a custom operation.
  op_type(str) - the operator type registered into quantizer. 
  The type should not conflict with pytorch_nndct

  attrs_list(Optional[List[str]], optional) - 
  the name list of attributes that define operation flavor. 
  For example, Convolution operation has such attributes as padding, dilation, stride and groups. 
  The order of name in attrs_list should be consistent with that of the arguments list. 
  Default: None


Perform the following steps:

  1. Aggregate some operations as a function. The first argument name of this function should be ctx. The meaning of ctx is the same as that in torch.autograd.Function
  2. Decorate this function with the decorator described above.

from pytorch_nndct.utils import register_custom_op

@register_custom_op(op_type="MyOp", attrs_list=["scale_1", "scale_2"])
def custom_op(ctx, x: torch.Tensor, y:torch.Tensor, scale_1:float, scale_2:float) -> torch.Tensor:
  return scale_1 * x + scale_2 * y 

class MyModule(torch.nn.Module):
  def __init__(self):

  def forward(self, x, y):
    return custom_op(x, y, scale_1=2.0, scale_2=1.0)
  1. Loop operation is not allowed in a custom operation.
  2. The number of return values for a custom operation can only be one.