文档

TensorFlow/Caffe Framework

You need to develop an operator plugin to interpret and map an operator in the TensorFlow or Caffe network to one adapted to the Ascend AI Processor.

This section uses the Add operator in the TBE operator development sample project as an example to describe how to implement the operator plugin.

MindStudio automatically generates the plugin code of the Add operator in the framework/tf_plugin/tensorflow_add_plugin.cc (use TensorFlow as an example) file.

  • Include the header file.
    1
2
    // Include the /compiler/include/register/register.h file in Ascend-CANN-Toolkit installation directory/ascend-toolkit/latest to use the operator registration class and call the operator registration API.
#include "register/register.h"
  • Register the plugin.
    1
2
3
4
5
6
7
    namespace domi{
REGISTER_CUSTOM_OP("Add")
    .FrameworkType(TENSORFLOW) 
    .OriginOpType({ge::AscendString("Add")})  
    .ParseParamsByOperatorFn(AutoMappingFn) 
    .ImplyType(ImplyType::TVM);    // Add it manually.
}
    • REGISTER_CUSTOM_OP: operator type registered with GE. According to Operator Analysis, the operator type is Add.
    • FrameworkType: TENSORFLOW indicates that the original framework type is TensorFlow, and CAFFE indicates that the original framework type is Caffe.
    • OriginOpType: type of the operator in the original framework.
    • ParseParamsByOperatorFn: function for registering models to be parsed. The AutoMappingFn function is used to automatically parse models.
    • ImplyType: operator implementation type, which needs to be manually added. ImplyType::TVM indicates a TBE operator and ImplyType::AI_CPU indicates an AI CPU operator.