文档首页 > > TBE自定义算子开发指南>

算子适配插件实现

算子适配插件实现

分享
更新时间:2020/07/23 GMT+08:00

本章节针对基于第三方框架,例如Tensorflow、Caffe等进行自定义算子开发的场景,开发人员完成自定义算子的实现代码后,需要进行适配插件的开发将基于第三方框架的算子映射成适配昇腾AI处理器的算子,将算子信息注册到GE中。基于第三方框架的网络运行时,首先会加载并调用GE中的插件信息,将第三方框架网络中的算子进行解析并映射成昇腾AI处理器中的算子。

若开发人员基于MindSpore框架进行自定义算子的开发,本章节跳过。

下面详细介绍如何进行算子插件的实现。

原理介绍

算子插件的实现包含昇腾AI处理器中算子类型的注册、原始框架中算子类型的注册以及原始框架中算子属性到昇腾AI处理器中算子属性的映射,算子的映射通过Parser模块完成。插件在整网络运行场景下的实现流程如图1所示。

图1 算子插件的实现流程
  1. 首先GE接收到第三方框架的原始网络模型,并进行初始化,网络模型的拓扑图我们简称为图。
  2. GE从Register注册模块中加载TBE算子插件生成的.so文件,在Host侧的存放路径为opp安装目录下的/usr/local/Ascend/ascend-toolkit/latest/xxx-linux_gccx.x.x/opp/framework/。
  3. 读取算子插件.so中的算子相关信息,并将其注册到算子插件的map文件中(所有算子插件的相关信息都会以map的形式存储到一个文件中)。
  4. GE向Parser模块发送调用Parser方法的请求。
  5. Parser模块根据算子类型(OpType)从算子插件的map文件中取出对应的Parser函数,并返回实现函数PaserParamsxx给Parser模块,Parser模块根据实现函数将第三方网络(Tensorflow/Caffe)算子中的属性映射到昇腾AI处理器中的算子属性,即算子原型中的属性定义,完成第三方网络中算子到昇腾AI处理器中算子的映射。
  6. 完成第三方网络图到适配昇腾AI处理器的图的转换后,后续会进行一些图融合、图优化的操作,然后最终将算子编译生成二进制文件。

插件代码实现

框架管理器(Framework)提供REGISTER_CUSTOM_OP宏,按照指定的算子名称完成算子的注册。

自定义算子的注册代码如下所示:
#include "register/register.h"
// #include "proto/caffe/caffe.pb.h"      //原始框架类型为Caffe时,需要包含此头文件
namespace domi
{
REGISTER_CUSTOM_OP("OpType")
    .FrameworkType(TENSORFLOW) 
    .OriginOpType("OriginOpType")
    .ParseParamsFn(ParseParamFunc)
    .ImplyType(ImplyType::TVM);
}
  • 在代码实现文件顶部使用预编译命令“#include”将插件实现函数相关的头文件包含到插件实现源文件中,头文件说明如下表所示。
    表1 头文件说明

    头文件

    目录

    作用

    register/register.h

    ATC安装目录下的/include/inc/register/register.h。

    包含该头文件,可使用算子注册类相关,调用算子注册相关的接口。

    proto/caffe/caffe.pb.h

    算子插件编译时依赖通过caffe.proto文件编译生成的caffe.pb.h文件。

    仅针对原始框架类型为Caffe时,需要包含此头文件;原始框架类型为Tensorflow时,不需要包含。

    算子插件编译时,会调用算子插件代码所在目录下的proto/caffe/caffe.pb.h文件进行算子参数解析。

  • REGISTER_CUSTOM_OP:注册自定义算子,OpType作为注册到GE中的算子类型,可以任意命名但不能和已有的算子命名冲突,且需要与3中的OpType保持一致。。
  • FrameworkType:TENSORFLOW代表原始框架为Tensorflow。CAFFE代表原始框架类型为Caffe。
  • OriginOpType:算子在原始框架中的类型。
  • ParseParamsFn:用来注册解析算子属性的函数。
    针对Tensorflow框架:
    • 若原始Tensorflow框架算子属性与昇腾AI处理器中算子属性一一对应(属性个数、属性名称与属性含义一致),可直接使用AutoMappingFn函数自动实现映射。
      .ParseParamsFn(AutoMappingFn) 

      AutoMappingFn函数会将Tensorflow框架中比昇腾AI处理器多的属性追加到昇腾AI处理器的算子属性

    • 若原始Tensorflow框架算子属性与昇腾AI处理器中算子属性无法对应,比如针对Conv2DBackpropInput算子,strides属性无法直接使用Tensorflow的对应算子中的strides属性,需要重新计算。所以需要在ParseParamsxx函数中进行对应的代码实现,如下所示:
      REGISTER_CUSTOM_OP("Conv2DBackpropInput")
          .FrameworkType(TENSORFLOW)
          .OriginOpType("Conv2DBackpropInput")
          .ParseParamsFn(ParseParamsConv2DBackpropInput)
          .ImplyType(ImplyType::TVM);
      } 
      // Conv2DBackpropInput算子的Parser函数实现
      Status ParseParamsConv2DBackpropInput(const Message* op_src, ge::Operator& op)
      {
          AutoMappingFn(op_src, op);  //调用AutoMapping函数完成昇腾AI处理器算子与Tensorflow算子之间可以对应起来的算子属性的映射
          ... ...
          SetStride(op);
          return SUCCESS;
      }
      //实现昇腾AI处理器的Conv2DBackpropInput算子中的Stride属性的赋值
      static bool SetStride(ge::Operator& op) {
          std::string dataFormat = "";
          int32_t hPosition = 0;
          int32_t wPosition = 0;
          if (ge::GRAPH_SUCCESS != op.GetAttr("data_format", dataFormat)){
              return false;
          }
          // get H & W position
          hPosition = dataFormat.find("H");
          wPosition = dataFormat.find("W");
          if(hPosition < 0 || wPosition < 0){
              return false;
          }
          std::vector<int32_t> strideList;
          if (GRAPH_SUCCESS == op.GetAttr("strides", strideList)) {
              if(strideList.empty())
              return false;
              }
              std::vector<int32_t> vec;
              vec.push_back(strideList[hPosition]);
              vec.push_back(strideList[wPosition]);
              op.SetAttr("strides", vec);
              return true;
          } else {
              return false;
          }
      }
    • 对于某些格式敏感的算子,需要在ParseParamsxx函数中将format设置为正确格式。

      由于Tensorflow原始图中不带format信息,GE在接收到Tensorflow原始图后会将图中所有格式设置为ND,但对于某些格式敏感算子,ND为不正确的格式,需要需要在插件的解析函数中进行format的设置。

      • 对于在Tensorflow的原始定义中有data format属性的算子,如果昇腾AI处理器的算子定义中输入或输出的格式与Tensorflow对应算子定义的data format一致,则format的设置AutoMapping函数会自动处理。
      • 对于在Tensorflow的原始定义中有data format属性的算子,如果昇腾AI处理器的算子定义中输入或输出的格式与Tensorflow对应算子定义的data format不一致,则需要在ParseParamsxx函数中将format设置为正确格式。
        例如针对Conv2D算子,原始定义中的data format为NCHW,但filter实际format为HWCN,此时需要在插件解析函数中将filter输入的format设置为HWCN。如下所示:
        const int POS_1 = 1;
        Status ParseParamsConv2D(const Message* op_src, ge::Operator& op) {
            AutoMappingFn(op_src, op);
            auto op_dsc = ge::OpDescUtils::GetOpDescFromOperator(op);
            ge::GeTensorDesc orgTensorW = op_dsc->GetInputDesc(POS_1);
            orgTensorW.SetOriginFormat(ge::FORMAT_HWCN);
            orgTensorW.SetFormat(ge::FORMAT_HWCN);
            auto ret = op_dsc->UpdateInputDesc(POS_1, orgTensorW);
            if (ret != ge::GRAPH_SUCCESS) {
                 return FAILED;
            }
            return SUCCESS;
        }
      • 对于在Tensorflow原始定义中没有data format属性的算子,例如若Tensorflow接口描述中只支持NHWC,则针对适配昇腾AI处理器的算子,则需要在插件的ParseParamsxx函数中强制设置输入输出format以及originFormat为NHWC。

        代码示例如下所示:

        Status ParseParamsSoftmaxMappingFn(const Message* op_src, ge::Operator& op) {
          AutoMappingFn(op_src, op);
          auto op_dsc = ge::OpDescUtils::GetOpDescFromOperator(op);
          ge::GeTensorDesc orgTensorW = op_dsc->GetInputDesc(0);
          ge::GeTensorDesc orgTensorW1 = op_dsc->GetOutputDesc(0);
          orgTensorW.SetOriginFormat(ge::FORMAT_NHWC);
          orgTensorW.SetFormat(ge::FORMAT_NHWC);
          orgTensorW1.SetOriginFormat(ge::FORMAT_NHWC);
          orgTensorW1.SetFormat(ge::FORMAT_NHWC);
          auto ret = op_dsc->UpdateInputDesc(0, orgTensorW);
          auto ret1 = op_dsc->UpdateOutputDesc(0, orgTensorW1);
          if(ret != ge::GRAPH_SUCCESS || ret1 != ge::GRAPH_SUCCESS) {
            return FAILED;
          }
          return SUCCESS;
        }

    针对Caffe框架,开发者需要自定义实现ParserParamFunc函数,实现算子属性的映射,函数实现方法请参见Caffe框架下算子解析函数实现

  • ImplyType:指定算子的实现方式,ImplyType::TVM表示该算子是TBE算子。

Caffe框架下算子解析函数实现

ParseParamsxx函数的声明如下所示:

Status ParseParamsxx(const Message* op_origin, ge::Operator& op_dest)
  • ParseParamsxx:函数名称,用户自定义,需要保持唯一。
  • op_origin:源算子模型,为protobuf格式的数据结构,来源于Caffe模型的proto文件,若用户自定义的算子在caffe.proto文件中未定义,则需要参考算子编译增加算子定义并编译生成caffe.pb.h与caffe.pb.cc文件。

    算子插件实现时会根据算子名称从caffe.proto编译生成的proto/caffe/caffe.pb.h文件与caffe.pb.cc文件中读取相关的算子属性进行解析,进而将算子数据结构转换为离线模型支持的数据结构。

  • op_dest:目标算子模型,适配昇腾AI处理器的离线模型的算子数据结构,保存算子信息,Operator类的详细描述请参见Operator类

下面详细介绍ParseParamsxx函数的具体实现。

  1. 定义指向LayerParameter的对象,并获取当前算子层的句柄。

    const caffe::LayerParameter* layer =dynamic_cast<const caffe::LayerParameter*>(op_origin);
    const caffe::xxxParameter& param = layer->reduction_param();

    其中:

    • param对象的类型caffe::xxxParameter中的xxxParameter需要与LayerParameter对象中声明的类型保持一致。
    • layer对象的成员函数xxx_param()的名字需要与与LayerParameter对象中声明的对象名字保持一致。

    例如caffe.proto中Reduction算子及Convolution算子的定义如下:

    message LayerParameter {
    optional ReductionParameter reduction_param = 136;
    optional ConvolutionParameter convolution_param = 106;
    ...
    }

    则获取当前Reduction算子层及Convolution算子层的句柄代码定义如下:

    const caffe::ReductionParameter& param = layer->reduction_param()

    const caffe::ConvolutionParameter& param = layer->convolution_param()

  2. 解析算子的每一个属性,并将其赋给Operator类型的对象op_dest。

    开发者可调用SetAtt接口将获取到的属性值赋给op_dest对象,SetAttr接口支持的数据类型及接口使用方法请参见SetAttr

    下面给出几种常见类型的参数解析示例:
    • int或者float类型参数

      例如,caffe.proto文件中算子参数定义如下所示:

      message ReductionParameter {
        ……
        optional int32 axis = 2 [default = 0];                                
        optional float coeff = 3 [default = 1.0]; 
      }

      调用SetAttr接口,将“param.axis()”的值赋给op_dest对象的axis属性,此处并将类型转换为int64_t,是为了提高中间计算数据的精度,由于原数据类型为int32,开发者也可以将数据类型转换为uint32_t;将“param.coeff()”的值赋给op_dest对象的coeff属性,coeff的数据类型为float,昇腾AI处理器中支持float类型的算子属性(支持的算子属性类型请参见SetAttr),所以不需要进行数据类型转换。代码示例如下所示:

      op_dest.SetAttr("axis", (int64_t)param.axis());
      op_dest.SetAttr("coeff", param.coeff());
    • enum类型参数

      例如,caffe.proto文件中算子属性定义如下所示:

      message ReductionParameter {
        enum ReductionOp {
          SUM = 1;
          ASUM = 2;
          SUMSQ = 3;
          MEAN = 4;
        }
      ...}
      1. 首先,需要将enum类型参数转换为map类型参数。
        std::map<caffe::ReductionParameter_ReductionOp, std::string> operation_map = {
                { caffe::ReductionParameter_ReductionOp_SUM, "SUM" },
            { caffe::ReductionParameter_ReductionOp_ASUM, "ASUM" },
            { caffe::ReductionParameter_ReductionOp_SUMSQ, "SUMSQ" },
            { caffe::ReductionParameter_ReductionOp_MEAN, "MEAN" },
            };
      2. 调用SetAttr接口,将map类型参数operation_map赋值给op_dest对象的operation属性。
        op_dest.SetAttr("operation", operation_map[param.operation()]);
    • repeated类型参数

      例如caffe.proto文件中算子参数定义如下所示:

      message xxxParameter {
      ...
        repeated float min_size = 1;
        repeated uint32 offset = 2;
      ....
      }
      1. 首先,将repeated float类型参数转换为vector<float>类型参数,将repeated uint32类型参数转换为vector<uint32_t>类型参数,并对vector类型的参数进行赋值。
        std::vector<float> min_size;      
        std::vector<uint32_t> offset;
        for(int i = 0; i < param.min_size_size(); ++i)
        {
           min_size.push_back(param.min_size(i));   //调用vector类型对象的push_back函数为min_size对象赋值
        }
        for(int i = 0; i < param.offset_size(); ++i)
        {
           offset.push_back(param.offset(i));      //调用vector对象的push_back函数为offset对象赋值
        }
      2. 调用SetAttr接口,将vector<float>类型参数min_size赋值给op_dest对象的min_size属性;将vector<uint32_t>类型参数offset赋值给op_dest对象的offset属性。
        op_dest.SetAttr("min_size", (min_size)); 
        op_dest.SetAttr("offset", (offset));

分享:

    相关文档

    相关产品

文档是否有解决您的问题?

提交成功!非常感谢您的反馈,我们会继续努力做到更好!
反馈提交失败,请稍后再试!

*必选

请至少选择或填写一项反馈信息

字符长度不能超过200

提交反馈 取消

如您有其它疑问,您也可以通过华为云社区论坛频道来与我们联系探讨

智能客服提问云社区提问