文档

ONNX Framework

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

MindStudio automatically generates the plugin code of the operator in the framework/onnx_plugin/xxx_plugin.cc 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
 8
 9
10
11
12
13
14
15
    using namespace ge;    // Add it manually.
namespace domi {
// Onnx ParseParams
Status ParseParamAdd(const Message* op_src, ge::Operator& op_dest) {
    // To do: Implement the operator plugin by referring to the Onnx Operator Development Guide.
    return SUCCESS;
}

// register op info to GE
REGISTER_CUSTOM_OP("Add")
    .FrameworkType(ONNX)   // Operator name with the original framework
    .OriginOpType("")      // Set the original frame type of the operator
    .ParseParamsByOperatorFn(ParseParamAdd)// Registering the callback function for parsing operator parameters 
    .ImplyType(ImplyType::TVM);    // Add it manually. TBE operator: ImplyType::TVM; AI CPU operator: ImplyType::AI_CPU
}  // namespace domi
    • REGISTER_CUSTOM_OP: registers a custom operator. Add is the operator type registered with GE. The value cannot conflict with existing operator names and must be the same as that registered in OpType.
    • FrameworkType: framework type. ONNX indicates that the original framework is ONNX.
    • OriginOpType: type of the operator in the original framework. Type it manually. For example, custom operator Add corresponds to the original framework type ai.onnx::11::Add of the ONNX OPP whose opset_version is 11. In this case, set this parameter to OriginOpType("ai.onnx::11::Add").
    • ParseParamsByOperatorFn(ParseParamAdd): registers a function for parsing operator attributes. You need to implement the ParseParamAdd callback function.

      The ParseParamAdd callback function is declared as follows:

      Status ParseParamAdd(const ge::Operator& op_src, ge::Operator& op_dest)
      • ParseParamAdd: function name, which is user-defined and must be unique.
      • op_src: an Operator class object defined by the ONNX framework, including attributes of the operator in the ONNX model. The definition is obtained from the original ONNX model file.
      • op_dest: CANN operator data structure, which stores the operator information.
    • 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.