更新时间:2021-03-18 GMT+08:00
分享

插件代码样例

  1. 修改“customop_te/plugin/caffe_reduction_layer.cpp”代码文件。

    • 修改“FilePath”为当前算子文件所在绝对路径+算子py文件名称。

      FilePath = "project_path/operator/reduction";

    • 修改FuncName为reduction.py文件中算子实现的函数名称。

      FuncName ="reduction"

    • 修改KernelName为reduction.py文件中设置的算子的KernelName。

      KernelName = "Reduction";

    • 修改算子的二进制文件*.o及描述文件.json所在的路径。

      te_bin_info.bin_file_path = "./kernel_meta/" + KernelName + ".o";

      te_bin_info.json_file_path = "./kernel_meta/" + KernelName + ".json";

      此路径为相对于当前工程目录的相对路径,例如Reduction.o文件在当前工程目录的operator/kernel_meta目录下。

  2. 修改后的reduction算子插件的代码样例如下所示:

    #include <Python.h>
    #include "custom/custom_op.h"
    #include "framework/omg/register.h"
    #include "framework/omg/omg_types.h"
    #include "proto/caffe/caffe.pb.h"
    #include "operator.h"
    #include "attr_value.h"
    #include <memory>
    #include <string>
    #include <vector>
    using namespace ge;
    
    namespace domi
    {
    // 用户自定义重写ParseParams函数,解析算子参数
    Status CaffeReductionParseParams(const Message* op_origin, ge::Operator& op_dest)
    {
        // 将op_origin转化为layer对象
        const caffe::LayerParameter* layer =
            dynamic_cast<const caffe::LayerParameter*>(op_origin);
    
        // 校验输入参数合法性
        if (nullptr == layer)
        {
            printf("Dynamic cast op_src to LayerParameter failed\n");
            return FAILED;
        }
        // 将算子参数中的枚举类型转换为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" },
        };
        // 获取当前算子层的句柄
        const caffe::ReductionParameter& param = layer->reduction_param();
        // 将算子中所有参数赋值给离线模型算子数据结构op_dest对象
        if(param.has_axis())
        {
            op_dest.SetAttr("axis", AttrValue::CreateFrom<AttrValue::INT>(param.axis()));
        }
        if(param.has_coeff())
        {
            op_dest.SetAttr("coeff", AttrValue::CreateFrom<AttrValue::FLOAT>(param.coeff()));
        }
        if(param.has_operation())
        {
            op_dest.SetAttr("operation", AttrValue::CreateFrom<AttrValue::STR>(operation_map[param.operation()]));
        }
        return SUCCESS;
    }
    
    // 自定义重写InferShapeAndType函数,获取算子的输出描述 
    Status CaffeReductionInferShapeAndType(const ge::Operator& op, vector<ge::TensorDesc>& v_output_desc)
    {
        // 获取算子输入的TensorDesc对象
        auto tensorDesc      = op.GetInputDesc(0);
        // 获取输入shape
        auto shape = tensorDesc.GetShape();
        int64_t axis = -1;
        
        ge::AttrValue axisAttrValue;
        if ((ge::GRAPH_SUCCESS != op.GetAttr("axis", axisAttrValue)) || (ge::GRAPH_SUCCESS != axisAttrValue.GetValue<AttrValue::INT>(axis)))
        {
            printf("Get axis failed!\n");
        }
        // In the OM model, all shape are supplemented to 4d. In this case, axis needs to be repaired to point to the original 2d.
        if (axis < 0) axis -= 2;
    
        if (axis < 0) axis += shape.GetDimNum();
    
        if (axis < 0 || axis >= shape.GetDimNum())
        {
            printf("invalid axis:%d, dim_size:%d\n", (int32_t)axis, (int32_t)shape.GetDimNum());
            return PARAM_INVALID;
        }
        int32_t dimsize = (int32_t)shape.GetDimNum();
        int32_t idx = 0;
        for(idx=axis; idx<dimsize; idx++)
        {
            shape.SetDim(idx, 1);
        }
        // 将调整后的shape信息设置到tensorDesc
        tensorDesc.SetShape(shape);    
        v_output_desc.push_back(tensorDesc);
    
        return SUCCESS;
    
    }
    
    
    // 编译算子,获取算子的二进制文件及算子描述文件
    Status CaffeReductionBuildTeBin(const ge::Operator& op, TEBinInfo& te_bin_info)
    {
        // 算子相关参数声明
        std::string FilePath   = "";
        std::string FuncName   = "";
        std::string KernelName = "";
       
        std::string operation  = "";
        int64_t     axis       = -1;
        float       coeff      = 1;
        // 获取算子参数值
        ge::AttrValue operationAttrValue;
        if ((ge::GRAPH_SUCCESS != op.GetAttr("operation", operationAttrValue)) || (ge::GRAPH_SUCCESS != operationAttrValue.GetValue<AttrValue::STR>(operation)))
        {
            // Add exception handling and maintenance information. 
            printf("GetOpAttr operation failed!\n");
        }
    
        // Parse the axis parameter. 
        ge::AttrValue axisAttrValue;
        if ((ge::GRAPH_SUCCESS != op.GetAttr("axis", axisAttrValue)) || (ge::GRAPH_SUCCESS != axisAttrValue.GetValue<AttrValue::INT>(axis)))
        {
            printf("GetOpAttr axis failed!\n");
        }
        // In the OM model, all shape are supplemented to 4d. In this case, axis needs to be repaired to point to the original 2d.
        if(axis < 0)
            axis -= 2;
    
        // Parse the coeff parameter. 
        ge::AttrValue coeffAttrValue;
        if ((ge::GRAPH_SUCCESS != op.GetAttr("coeff", coeffAttrValue)) || (ge::GRAPH_SUCCESS != coeffAttrValue.GetValue<AttrValue::FLOAT>(coeff)))
        {
            printf("GetOpAttr coeff failed!\n");
        }
        // 获取算子输入描述信息
        TensorDesc input_desc      = op.GetInputDesc(0);
    
        // 解析Input shape并校验算子的Dim是否为4
        if(input_desc.GetShape().GetDimNum() != 4)
        {
            printf("The shape size is %d, which is not 4!", (int32_t)input_desc.GetShape().GetDimNum());
            return FAILED;
        }
        // 设置算子文件路径
        FilePath   = "/home/ascend/tools/projects/customop_te/operator/reduction";
        // 设置实现算子的函数名
        FuncName   = "reduction";
        // 设置算子实现文件中定义的算子名称
        KernelName = "Reduction";
    
        // i => int; s => string; f => dobule; O => bool, and bool value is Py_True or Py_False,调用Tensor Engine的编译接口编译算子
        te::BuildTeCustomOp(te_bin_info.ddk_version, op.GetName(), FilePath, FuncName,
                        "(i,i,i,i), s, i, s, f, s", input_desc.GetShape().GetDim(0), input_desc.GetShape().GetDim(1),
                        input_desc.GetShape().GetDim(2), input_desc.GetShape().GetDim(3), "float16", axis, operation.c_str(), coeff,
                        KernelName.c_str());    //此处的参数顺序需要与算子实现函数中参数顺序保持一致。
    
        // set te op json to te_bin_info 
        te_bin_info.bin_file_path  = "./kernel_meta/" + KernelName + ".o";
        te_bin_info.json_file_path = "./kernel_meta/" + KernelName + ".json";
    
        return SUCCESS;
    }
    
    REGISTER_CUSTOM_OP("custom_reduction") //####test_reduction is the type name of the operator in the OM model. It can be specified randomly and cannot be the same as an existing type name. It is case sensitive. 
        .FrameworkType(CAFFE)  // Enumerated type. The options are as follows: CAFFE, TENSORFLOW
        .OriginOpType("Reduction")  // Reduction indicates the type name of the operator in the caffe framework.
        .ParseParamsFn(CaffeReductionParseParams)  // Op parameters parse function
        .InferShapeAndTypeFn(CaffeReductionInferShapeAndType)       // Set output description and datatype function
        .TEBinBuildFn(CaffeReductionBuildTeBin)           // Build Te op binary function
        .ImplyType(ImplyType::TVM);        // Implementation type. Enumerated type, The options are as follows: TVM, AI_CPU.
    
    }  // namespace domi

分享:

    相关文档

    相关产品