TensorFlow/Caffe框架
开发者需要进行算子适配插件的开发,实现将TensorFlow/Caffe框架网络中的算子进行解析并映射成昇腾AI处理器中的算子。
本章节将以TBE算子开发样例工程中“Add”算子为例讲解算子适配插件的实现,供开发者参考。
MindStudio在“framework/tf_plugin/tensorflow_add_plugin.cc”(以TensorFlow为例)文件已自动生成了Add算子的插件代码。
- 包含头文件。
1 2
//包含该头文件,可使用算子注册类相关,调用算子注册相关的接口,为Ascend-cann-toolkit安装目录/ascend-toolkit/latest/compiler/include/register/register.h文件 #include "register/register.h"
- 进行插件注册。
1 2 3 4 5 6 7
namespace domi{ REGISTER_CUSTOM_OP("Add") .FrameworkType(TENSORFLOW) .OriginOpType({ge::AscendString("Add")}) .ParseParamsByOperatorFn(AutoMappingFn) .ImplyType(ImplyType::TVM); //需手动添加 }
- REGISTER_CUSTOM_OP:算子注册到GE的算子类型,根据算子分析,算子类型为“Add”。
- FrameworkType:TENSORFLOW代表原始框架类型为TensorFlow,CAFFE代表原始框架为Caffe。
- OriginOpType:算子在原始框架中的类型。
- ParseParamsByOperatorFn:用来注册解析模型的函数,使用AutoMappingFn函数自动实现解析。
- ImplyType:指定算子的实现方式,需要手动添加。ImplyType::TVM表示该算子是TBE算子,ImplyType::AI_CPU表示该算子是AI CPU算子。
父主题: 算子适配插件实现