Atlas 500应用-算子注册
算子注册
框架管理器(Framework)提供REGISTER_CUSTOM_OP宏,按照指定的算子名称完成算子的注册。
自定义算子的注册代码如下所示:
REGISTER_CUSTOM_OP("test_layer") .FrameworkType(CAFFE) .OriginOpType("Test") .ParseParamsFn(ParseParamsxx) .InferShapeAndTypeFn(InferShapeAndTypexx) .TEBinBuildFn(BuildTeBinxx) .ImplyType(ImplyType::TVM) .Formats({DOMI_TENSOR_NC1HWC0}, {DOMI_TENSOR_NC1HWC0}) .WeightFormats({DOMI_TENSOR_FRACTAL_Z, DOMI_TENSOR_NC1HWC0});
其中:
- REGISTER_CUSTOM_OP:注册自定义算子,"test_layer"作为离线模型文件中的算子名称,可以任意命名但不能和已有的算子命名冲突。
- FrameworkType:不同框架的算子参数解析逻辑不同,因此对于不同框架模型需要有不同的插件,在插件的注册代码需要表明对应的框架。当前请配置为CAFFE。
- OriginOpType:算子类型,需要与Caffe Prototxt中定义的算子类型保持一致,否则无法正常解析,内置的caffe.proto文件路径为DDK安装路径下的“/include/inc/custom/proto/caffe/caffe.proto”。
- ParseParamsFn:用来注册解析模型的函数,“ParseParamsxx”即解析算子中实现的函数,针对Caffe框架开发插件时需要,若自定义的算子为重写昇腾AI处理器已经支持的算子,则此步骤可跳过;若自定义的算子为昇腾AI处理器不支持的算子,则此步骤必选。
- InferShapeAndTypeFn:用来注册形状和类别推断函数,“InferShapeAndTypexx”即推理算子输出张量描述中实现的函数。
- TEBinBuildFn:用来注册TE算子编译函数,“BuildTeBinxx”即编译算子中实现的函数。
- ImplyType:指定算子的实现方式,ImplyType::TVM表示该算子是TE算子。
- Formats:算子输入数据与输出数据的数据排布格式,其中第一个列表是输入数据格式列表,第二个列表为输出数据格式列表,若输入数据有多个,请在第一个列表中列出每一个输入数据的排布格式。例如若有两个输入数据,分别是NC1HWC0格式,则Formats函数调用方式如下:
.Formats({DOMI_TENSOR_NC1HWC0, DOMI_TENSOR_NC1HWC0}, {DOMI_TENSOR_NC1HWC0})
详细说明请参见《Framework API参考》中的“Formats函数”。
- WeightFormats:设置算子的权重数据的排布格式,支持的数据格式类型请参见《Framework API参考》中的“WeightFormats函数”。例如Convolution的filter的数据排布格式是fractal_Z,bias的数据排布格式是NC1HWC0。
如果模型转换时开启量化开关,需要在此接口中增加Framework处理时新增的常量数据的格式,当前Framework支持的量化算子有Conv,FC与Depthwise Conv,若这几个算子在模型转换时开启了量化开关,则需要在WeightFormats接口的参数列表最后增加6个DOMI_TENSOR_NC1HWC0的数据格式(量化时Framework新增了6个数据排布格式为NC1HWC0的常量,此处不详细介绍),开启量化后的Convolution的WeightFormats接口示例如下:
.WeightFormats({DOMI_TENSOR_FRACTAL_Z, DOMI_TENSOR_NC1HWC0, DOMI_TENSOR_NC1HWC0, DOMI_TENSOR_NC1HWC0, DOMI_TENSOR_NC1HWC0,DOMI_TENSOR_NC1HWC0, DOMI_TENSOR_NC1HWC0, DOMI_TENSOR_NC1HWC0})