要将量化模型转换为 XMODEL、vai_q_pytorch 会提供修饰器来允许您将单一运算或一组运算寄存为 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 names in attrs_list should be consistent with that of the arguments list.
Default: None
"""
执行以下步骤:
- 将一些运算聚合为一个函数。该函数的第一个实参名应为 ctx,其含义与 torch.autograd.Function 中的含义相同
- 使用前述修饰器来修饰此函数。
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)
限制:- 自定义运算中不允许循环运算。
- 自定义运算只能返回一个值。