文档
注册

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算子。
搜索结果
找到“0”个结果

当前产品无相关内容

未找到相关内容,请尝试其他搜索词