algorithm_register

Applicability

Product	Supported
Atlas 350 Accelerator Card	√
Atlas A3 training product/Atlas A3 inference product	√
Atlas A2 training product/Atlas A2 inference product	√
Atlas 200I/500 A2 inference product	x
Atlas inference product	x
Atlas training product	x

Description

Registers the custom algorithm provided by the user with AMCT.

Prototype

algorithm_register(name, src_op, quant_op, deploy_op)

Parameters

Parameter	Input/Output	Description
name	Input	Algorithm name. Data type: string
src_op	Input	Operator to be replaced. Data type: string
quant_op	Input	Quantized operator. Data type: torch.nn.Module
deploy_op	Input	Deployable operator. Data type: torch.nn.Module

Returns

None

Restrictions

None

Example

# Custom algorithm name
name = 'customize_algo'
# Type of the operator to be quantized
src_op = 'Linear'
# Quantization operator implemented by the user
class CustomizedQuantOp(BaseQuantizeModule):
    def __init__(self,
                 ori_module,
                 layer_name,
                 quant_config):
        super().__init__(ori_module, layer_name, quant_config)

    @torch.no_grad()
    def forward(self, inputs):
        return
quant_op = CustomizedQuantOp
# Deployable operator implemented by the user
class CustomizedDeployOp(torch.nn.Module):
    def __init__(self, quant_module):
        super().__init__()

    def forward(self, x):
        return
deploy_op = CustomizedDeployOp
# Register the custom algorithm.
algorithm_register(name, src_op, quant_op, deploy_op)

Parent topic: Torch Module-based Quantization APIs