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. Args: 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:
- 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
- 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)
- Loop operation is not allowed in a custom operation.
- The number of return values for a custom operation can only be one.