适配插件开发(Caffe框架)

简介

您可以参考本章节进行算子适配插件的开发,将基于第三方框架的算子映射成适配昇腾AI处理器的算子,将算子信息注册到Graph Engine(简称:GE)中。基于Caffe框架的网络运行时,首先会加载并调用GE中的插件信息,将原始框架网络中的算子进行解析并映射成适配昇腾AI处理器中的算子。

下文我们将适配昇腾AI处理器的算子称为CANN算子。

Caffe框架算子到CANN算子的一对多、多对多、多对一映射暂不支持。

原理介绍

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

图1 算子插件的实现流程
  1. 首先GE接收到第三方框架的原始网络模型,并进行初始化,网络模型的拓扑图我们简称为图。
  2. GE从Register注册模块中加载算子插件生成的.so文件,在CANN软件安装后文件存储路径“opp/built-in/framework/”路径中。
  3. 读取算子插件.so中的算子相关信息,并将其注册到算子插件的map文件中(所有算子插件的相关信息都会以map的形式存储到一个文件中)。
  4. GE向Parser模块发送调用Parser方法的请求。
  5. Parser模块根据算子类型(OpType)从算子插件的map文件中取出对应的Parser函数,并返回实现函数ParseParamsByOperatorFn给Parser模块,Parser模块根据实现函数将第三方网络算子中的属性映射到CANN算子的属性,即算子原型中的属性定义,从而完成第三方网络中算子到CANN算子的映射。
  6. 后续会进行图准备、图拆分及图优化等一系列操作,最终生成适配昇腾AI处理器的网络模型。

插件实现

GE提供REGISTER_CUSTOM_OP宏,按照指定的算子名称完成算子的注册。

原始框架为Caffe的自定义算子注册代码如下所示:
#include "register/register.h"
#include "graph/operator.h"
namespace domi
{
REGISTER_CUSTOM_OP("OpType")
    .FrameworkType(CAFFE) 
    .OriginOpType("OriginOpType")
    .ParseParamsByOperatorFn(ParseParamByOpFunc)   // 用来注册解析算子属性的函数
    .ImplyType(ImplyType::TVM);    // TBE算子:ImplyType::TVM;AI CPU算子:ImplyType::AI_CPU
}