カスタム演算 (OP) の登録 - 2.5 日本語

Vitis AI ユーザー ガイド (UG1414)

Document ID
UG1414
Release Date
2022-06-15
Version
2.5 日本語
量子化済みモデルを xmodel に変換するため、vai_q_pytorch には 1 つの演算 (op) または op グループを XIR にとって未知のカスタム op として登録するデコレーターが用意されています。
# 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

  """

次の手順を実行します。

  1. いくつかの op を 1 つの関数にまとめます。この関数の最初の引数名は ctx とします。ctx の意味は、torch.autograd.Function のものと同じです。
  2. この関数を、上記のデコレーターで修飾します。

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. カスタム op では、ループ演算はサポートされません。
  2. カスタム op は戻り値を 1 つしかとることができません。