algorithm_register
【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct
产品支持情况
| 产品 | 是否支持 |
|---|---|
| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | √ |
| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ |
功能说明
将用户提供的自定义算法注册到AMCT工具。
函数原型
algorithm_register(name, src_op, quant_op, deploy_op)参数说明
| 参数名 | 输入/输出 | 说明 |
|---|---|---|
| name | 输入 | 含义：算法名称。 数据类型：string。 |
| src_op | 输入 | 含义：替换的算子。 数据类型：string。 |
| quant_op | 输入 | 含义：量化算子。 数据类型：torch.nn.Module。 |
| deploy_op | 输入 | 含义：部署算子。 数据类型：torch.nn.Module。 |
返回值说明
无
调用示例
# 自定义算法名称 name = 'customize_algo' # 需要量化的算子类型 src_op = 'Linear' # 用户自己实现的量化算子 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 # 用户自己实现的部署算子 class CustomizedDeployOp(torch.nn.Module): def __init__(self, quant_module): super().__init__() def forward(self, x): return deploy_op = CustomizedDeployOp # 注册自定义算法 algorithm_register(name, src_op, quant_op, deploy_op)【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
