自定义算子开发

简介

地平线工具链中已经支持了丰富的算子,在大多数情况下,您的模型应该可以通过前文所述模型转换顺利部署到地平线计算平台上。已支持的算子情况可以参考 工具链算子支持约束列表-ONNX Operator Support List 章节。少部分算子不支持情况下,我们强烈建议您先尝试下替换算子的可能性,这样有利于将地平线计算平台能力充分发挥出来,且开发成本会更低。

自定义算子只提供CPU上算子开发能力,一个完整的自定义算子应用过程包括创建模板、算子实现、算子编译、含自定义算子模型转换和运行含自定义op模型几个阶段。具体流程如下图所示:

如图所示,定义自定义OP需要有两部分的任务:在模型转换阶段,需要提供自定义OP的python代码;在模拟器/上板运行推理阶段,需要提供自定义OP的C++代码。要确保这两部分的代码运算的一致性。

环境配置

在进行含自定义算子的模型转换之前,需要先对环境进行配置,请务必保证您的环境满足本节内描述的要求。

注意

请注意,如需在X86环境下对含自定义算子的模型进行编译或推理,需保证gcc版本大于5.0。

请参考下方,对 环境变量 进行配置:

  1. GCC环境配置

    x86_gcc_path="/bin/g++-11" x86_gcc_lib="/usr/lib/gcc/x86_64-linux-gnu/11/" cross_gcc_root="/usr/" cross_gcc_aarch64_lib="${cross_gcc_root}/aarch64-linux-gnu/lib/" cross_gcc_x86_lib="${cross_gcc_root}/lib/x86_64-linux-gnu/"
  2. 编译器依赖配置

    hbdk_path="$(pip show hbdk4-compiler |grep Location |awk -F " " '{printf$2}')/hbdk4" hbtl_aarch64="${hbdk_path}/runtime/aarch64_unknown_linux_gnu/nash" hbtl_x86_64="${hbdk_path}/runtime/x86_64_unknown_linux_gnu/nash" export INTERPRETER_ENABLE=0
  3. X86环境变量设置

    export HBDK_TARGET_X86_64_UNKNOWN_LINUX_GNU_CXX="${x86_gcc_path}" export HBDK_TARGET_X86_64_UNKNOWN_LINUX_GNU_CXXFLAGS="-I${hbtl_x86_64}/include/ -L${hbtl_x86_64}/lib -L${hbdk4_path}/compiler -L${hbdk4_path}/compiler/_mlir_libs -std=c++14" export HBDK_TARGET_X86_64_UNKNOWN_LINUX_GNU_LDFLAGS="-L${hbtl_x86_64}/lib -static-libstdc++ -L${x86_gcc_lib} -Wl,-rpath,${x86_gcc_lib}"
  4. Arm环境变量设置

    export HBDK_TARGET_AARCH64_UNKNOWN_LINUX_GNU_CXX="${cross_gcc_root}/bin/aarch64-linux-gnu-g++" export HBDK_TARGET_AARCH64_UNKNOWN_LINUX_GNU_CXXFLAGS="-I${hbtl_aarch64}/include/ -L${hbtl_aarch64}/lib -L${cross_gcc_aarch64_lib} -L${hbdk4_path}/compiler/_mlir_libs -L${hbdk4_path}/compiler/_mlir_libs -v" export HBDK_TARGET_AARCH64_UNKNOWN_LINUX_GNU_LDFLAGS="-L${hbtl_aarch64}/lib -L${cross_gcc_aarch64_lib} -static-libstdc++ -Wl,-rpath,${cross_gcc_aarch64_lib}"
  5. 自定义算子动态库关联

    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:"${x86_gcc_lib}:${cross_gcc_x86_lib}"

自定义算子(中间节点)的模型转换

自定义算子(中间节点)是指模型中某个/某些算子注册为自定义算子的场景。

模型文件修改

在准备好自定义算子实现后,为了将算子应用起来,您需要从原始模型文件和模型转换配置两个方面做出相应调整(下面分别以 Caffe 模型和 ONNX 模型为例)。

Caffe 模型

原始模型文件中,将自定义算子对应的算子类型标记为 Custom,并提供一组 custom_param,示例如下。

layer { name: "hr_op" type: "Custom" bottom: "res3d_in" top: "res3d" custom_param { kind: "CustomIdentity" shape { dim: 1 dim: 512 dim: 28 dim: 28 } params: "'kernel_size': 10 \n'threshold': 0.5" } }

以上完整 custom_param 示例中:

  • kind 是自定义算子的内部实现名称,该自定义OP为恒等OP,因此命名为 CustomIdentity,该名称在后续Python及C++代码中均会体现。

  • shape 是算子的输出尺寸,需要完整指定。

  • params 是算子的传入参数指定形式为 'param_name': param_value,多个参数之间使用 \n 分隔。

在模型转换配置中,使用自定义算子需要在配置文件中加入一个新的自定义op参数组如下:

#... custom_op: # 自定义op的校准方式 custom_op_method: register # 自定义OP的实现文件 op_register_files: sample_custom.py

对于 Caffe 模型,以上参数组中的两个参数都是必须配置的。custom_op_method 固定使用 registerop_register_files 是自定义算子计算的实现文件,请使用相对路径。

完成这些配置后,模型转换的后续步骤与其他一般模型转换过程一致。

ONNX 模型

  1. 含有自定义算子的Onnx模型的获取:
  • 从pytorch等其他框架转换而来

    import torch from horizon_nn.horizon_onnx.onnx_pb import TensorProto from torch.onnx.symbolic_helper import parse_args from torch.onnx.utils import register_custom_op_symbolic from torch import Tensor model = torch.hub.load('pytorch/vision:v0.10.0', 'googlenet', pretrained=True) def _transform_input(x: Tensor) -> Tensor: return x model._transform_input = _transform_input @parse_args("v", "v") def horizon_pool(g, input, output_size): return g.op( 'horizon.custom::PyOp', #required, ! must be 'horizon.custom' domain ! input, class_name_s="GlobalAveragePool", #required ! must match the class def name in sample_custom python file ! compute_s="compute", #optional, 'compute' by default module_s="sample_custom", #required ! must match the file name of the "op_register_files" ! input_types_i=[TensorProto.FLOAT], #required output_types_i=[TensorProto.FLOAT], #required output_shape_s=["1, 1024, 1, 1"]) #required d_input = torch.rand(1, 3, 224, 224) register_custom_op_symbolic('::adaptive_avg_pool2d', horizon_pool, opset_version=11) torch.onnx.export(model, d_input, "googlenet_cop.onnx", opset_version=11)
  • 直接生成onnx模型

    参考代码:

    import onnx import numpy as np from onnx import helper, checker, shape_inference, numpy_helper, TensorProto def make_normal_data(shape): return np.random.normal(loc=0.0, scale=1.0, size=shape).astype(np.float32) # conv def make_simple_model(): # create nodes conv_input_shape = (1, 3, 224, 224) conv_output_shape = (1, 3, 224, 224) add_param_shape = (1, 3, 224, 224) add_1_param_data = np.zeros(add_param_shape).astype(np.float32) add_2_param_data = np.ones(add_param_shape).astype(np.float32) conv_weight_shape = (3, 3, 3, 3) conv_output_shape = (1, 3, 224, 224) conv_weight_data = make_normal_data(conv_weight_shape) add_1_node = helper.make_node( "PyOp", # required, 类型必须是'PyOp' name="add_1", # required, 不同的op名称不能相同 inputs=["input0", "add_1_param"], # required, 需要是一个list, 且需要与实现文件中的输入数量保持一致 outputs=["add_1_out"], # required, 需要是一个list, 且需要与实现文件中的输出数量保持一致 domain="horizon.cop1", # required, 不同实现逻辑的自定义算子实现需要通过不同的domain名称来实现 class_name="Cop1", # required, 需要与自定义算子的实现文件中的class名称一致 module="custom_op.horizon_ops", # required, 需要与包含自定义算子的实现文件的路径一致 compute="compute", # required, 需要与自定义算子实现class中的计算逻辑函数一致 ext_compute = "func1", # 配置numba时需要配置此选项,用于指定numba的函数名称 input_types=[ TensorProto.FLOAT, TensorProto.FLOAT, ], # required, 需要是一个list, 其长度需要与该算子的inputs属性数量一致, 且与实现文件中的输入数量保持一致 output_types=[ TensorProto.FLOAT ], # required, 需要是一个list, 其长度需要与该算子的outputs属性数量一致, 且与实现文件中的输出数量保持一致 output_shape=["1, 3, 224, 224"], # optional, 在模型中未添加pyop的输出 value_info时, 必须填写 ) add_2_node = helper.make_node( "PyOp", name="add_2", inputs=["input1", "add_1_out", "add_2_param"], outputs=["add_2_out", "output0"], domain="horizon.cop2", class_name="Cop2", module="custom_op.horizon_ops", compute='compute', input_types=[TensorProto.FLOAT, TensorProto.FLOAT, TensorProto.FLOAT], #required output_types=[TensorProto.FLOAT, TensorProto.FLOAT], #required output_shape=["1, 3, 224, 224", "1, 3, 224, 224"]) conv_1_node = helper.make_node("Conv", inputs=["add_2_out", "W0"], outputs=["output1"], dilations=(1, 1), group=1, kernel_shape=(3, 3), pads=(1, 1, 1, 1), name="conv_1") # nodes nodes = [add_1_node, add_2_node, conv_1_node] # inputs model_input_1 = helper.make_tensor_value_info("input0", TensorProto.FLOAT, conv_input_shape) model_input_2 = helper.make_tensor_value_info("input1", TensorProto.FLOAT, conv_input_shape) # Outputs model_output_1 = helper.make_tensor_value_info("output0", TensorProto.FLOAT, conv_output_shape) model_output_2 = helper.make_tensor_value_info("output1", TensorProto.FLOAT, conv_output_shape) # Intermediate tensors add_1_out = helper.make_tensor_value_info("add_1_out", TensorProto.FLOAT, conv_output_shape) add_2_out = helper.make_tensor_value_info("add_2_out", TensorProto.FLOAT, conv_output_shape) # create constant tensor W0_tensor = helper.make_tensor("W0", TensorProto.FLOAT, conv_weight_shape, conv_weight_data.flatten()) add_1_param = helper.make_tensor("add_1_param", TensorProto.FLOAT, add_param_shape, add_1_param_data.flatten()) add_2_param = helper.make_tensor("add_2_param", TensorProto.FLOAT, add_param_shape, add_2_param_data.flatten()) # make graph graph = helper.make_graph( nodes, "simple_conv_model", inputs=[model_input_1, model_input_2], # input outputs=[model_output_1, model_output_2], # output initializer=[W0_tensor, add_1_param, add_2_param], # initializer value_info=[add_1_out, add_2_out], # value_info ) # make model onnx_model = helper.make_model(graph, opset_imports=[ helper.make_opsetid("", 11), helper.make_opsetid("horizon.cop1", 1), helper.make_opsetid("horizon.cop2", 1) ], producer_name="onnx-test") # shape inference onnx_model = shape_inference.infer_shapes(onnx_model) # # model check checker.check_model(onnx_model) # save model onnx.save(onnx_model, "custom_op.onnx")
注意

Onnx模型中PyOp属性的注意点:

  • domain属性一定要设置,不然的话会被默认成onnx标准domain从而报错。不同实现的自定义算子需要设置在不同的domain下。

  • module需要与注册时使用的注册文件同名。若注册文件在当前目录的子文件夹中,则需要修改module内容。例如:若 sample_custom.py 在当前路径的custom_op 文件夹中,则该module应设置为 custom_op.sample_custom

  • 目前仅onnx模型支持多种类型的自定义算子,如您需要在其他框架中支持多种类型的自定义算子请联系地平线技术支持人员。

  1. Caffe 模型一致,需要在模型转换配置中,使用自定义算子需要在配置文件中加入一个新的自定义op参数组如下:

    #... custom_op: # 自定义op的校准方式 custom_op_method: register # 自定义OP的实现文件 op_register_files: sample_custom.py

    对于 ONNX 模型,以上参数组中的两个参数都是必须配置的。custom_op_method 固定使用 registerop_register_files 是自定义算子计算的实现文件,请使用相对路径。

完成这些配置后,模型转换的后续步骤与其他一般模型转换过程一致。

算子实现

在模型转换阶段,需要提供自定义算子的Python实现,工具会利用该实现函数完成模型校准必需的推理阶段。

注意

请注意,由于工具在PTQ转换过程中会以 working_dir 为工作目录,我们强烈建议算子实现中涉及工作目录的配置时,配置为绝对路径,如需要配置为相对路径,请以 working_dir 为工作目录进行相对路径的指定。

Python模板文件(sample_custom.py)如下:

from horizon_nn.custom.op_registration import op_implement_register, op_shape_infer_register @op_implement_register("CustomIdentity") class CustomIdentity(object): def __init__(self, kernel_size, threshold): self._kernel_size = kernel_size self._default_threshold = threshold def compute(self, X): return X @op_shape_infer_register("CustomIdentity") def infer_shape(inputs_shape): outputs_shape = inputs_shape return outputs_shape

custom_op示例中的配置文件(horizon_ops.py)如下:

from horizon_nn.custom.op_registration import op_implement_register import numba @numba.njit def func1(x1, x2): out = x1 + x2 return out @numba.njit def func2(x1, x2, x3): out = x1 + x2 + x3 + 2 return out, out @op_implement_register("Cop1") class Cop1(object): def __init__(self, ext_compute): self.func1 = ext_compute def compute(self, x1, x2): return eval(f"{self.func1}(x1, x2)") @op_implement_register("Cop2") class Cop2(object): def __init__(self, ext_compute): self.func2 = ext_compute def compute(self, x1, x2, x3): return eval(f"{self.func2}(x1, x2, x3)")

该文件的名字(sample_custom.py)需要填入模型转换的yaml配置文件中 op_register_files ,否则工具无法正常import自定义算子定义, 并且修饰器 op_implement_register 注册的custom op类的名称 CustomIdentity 需要与Caffe自定义OP的属性 kind 或者Onnx自定义OP的属性 class_name 一致。

对于 Caffe 模型, init 函数中的参数(kernel_size, threshold)都是通过prototxt文件中的 params 传入的, 用于自定义op模块的初始化。 op_shape_infer_register 用于Caffe模型的算子shape注册。

对于 Onnx 模型,自定义op的shape解析有两种方式,可以通过在创建onnx模型时,将pyop输出的value_info添加到onnx模型中,或者在对应的pyop中创建output_shape属性。 还需要注意自定义算子中的 module 必须与存放自定义算子实现的文件保持一致, 如果属性设置为 custom_op.horizon_ops , 则自定义算子实现的文件名称为 horizon_ops ,且要放在 custom_op文件夹中, 保持与onnx模型的层级关系。由于同一个domain中的同名算子实现必须相同, 因此不同的自定义算子的 domain 属性需要不同。

上述操作完成后即可进行浮点转定点的操作,得到相应的hbm文件及用于含自定义算子模型部署的动态库.so文件。

自定义算子(前后处理)的模型转换

自定义算子(前后处理)是指通过添加自定义算子的方式将前/后处理逻辑集成到模型中的场景。

为降低含自定义算子模型模型板端部署的难度,同时提升前后处理代码复用性,我们也提供了这种场景。该场景的使用会涉及环境变量的设置及前后处理代码的编写,您可参考下方示例。

注意

请注意,在进行下方操作前,请确保已按照上文 环境配置 章节完成环境变量的配置。

前后处理示例说明

下方示例代码,串联GoogleNet模型,将前后处理集成到HBIR模型中。

import os import sys import numpy as np from numba import njit, float32 from hbdk4.compiler import translate, compile, load, save, leap, convert import hbdk4 # 前处理numba实现 @njit def preprocess(data): ... return preprocessed_data # 后处理numba实现 @njit def postprocess(result): ... return res module = load("googlenet_224x224_yuv444_quantize.bc") # pipeline链路 def pipeline(image): # 前处理 preprocessed_data = preprocess(image) # 调度模型推理 infer_result = leap.call(module, "googlenet_224x224_yuv444", preprocessed_data) # 后处理 res = postprocess(infer_result) return res # 输入数据,用于串联并验证链路 random_data = np.random.randn(1, 3, 224, 224) data = (random_data * 127.5 + 128).astype(np.uint8) # 将前后处理集成到HBIR模型中 module = translate(pipeline, data) # HBIR模型量化 b25_module = convert(module, "b25") #save(b25_module, "pipeline.bc") # 重新编译HBIR模型,生成新的HBM compile(b25_module, "googlenet_pipeline.hbm", "b25", opt=0, progress_bar=True)

产出物说明

在完成上述操作后,您将得到一系列产出物文件,其中需要重点关注的文件包括:

  • libpipeline_xxxxxxxx_arm.so:Arm自定义算子动态库,用于模型在板端推理时使用。

  • *.hbm:含前/后处理的用于上板的离线模型文件。

含自定义算子的模型上板运行

在拿到hbm文件以及自定义算子运行相关的动态库.so文件后,还不能直接在开发板上运行,如想要上板运行,还需要配置对应动态库的加载路径,配置方法如下:

# 创建文件夹存放与自定义算子运行相关的动态链接库, xxx需要替换为您的真实路径(绝对路径) mkdir -p xxx/custom_libs cp xxx.so xxx/custom_libs # 配置环境变量 export HB_DNN_CUSTOM_HBTL_PATH=xxx/custom_libs