产品文档6.5. 自定义算子开发

6.5. 自定义算子开发

2026-02-06 20:50:11

6.5.1. 简介

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

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

../../../_images/custom_op_development.png

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

 

6.5.2. 含自定义算子的模型转换

6.5.2.1. 模型文件修改

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

6.5.2.1.1. Caffe 模型¶

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

以上完整 custom_param 示例中:

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

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

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

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

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

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

6.5.2.1.2. ONNX 模型

1.含有自定义算子的Onnx模型的获取:

  • 从pytorch等其他框架转换而来

  • 直接生成onnx模型

参考代码:

注意

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

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

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

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

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

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

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

6.5.2.2. 算子实现

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

注意

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

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

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

该文件的名字(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 属性需要不同。

上述操作完成后即可进行浮点转定点的操作,得到相应的bin文件。

 

6.5.3. 含自定义算子的上板运行

在拿到bin文件后,还不能直接在开发板上运行。在运行之前需要先提供自定义算子的C++代码实现。 您可以使用下文提供的模板进行修改并添加到示例代码中进行使用。

如果您只是希望测试自定义算子的功能,也可以直接使用我们提供的模板文件,模板文件中将输入直接赋值为输出使用, 所以这个自定义算子并不会对结果造成任何影响。

6.5.3.1. 自定义算子C++模板

Runtime模板文件如下:

注解

该函数名称的前缀(即 Cop1 和 Cop2) 需要与自定义OP的类型( Kind )一致, 其传入的参数为:

  • bottom_blobs :自定义OP节点输入数据。

  • top_blobs :自定义OP节点输出数据。

  • inferCtrlParam :自定义算子初始化阶段的输入参数。

注意

模板中的运算规则均为输出等于所有输入数据累加后再加上数值1,因此后续若需要定义其他行为,则需相应的更改运算规则即可。

6.5.3.2. 自定义算子注册

当您完成C++版本模板的修改后,仅需要在示例的CMakeLists.txt中添加对模板文件的包含, 并在应用程序中增加对算子的注册即可,注册请参考以下代码:

当您完成对模板文件的依赖及算子注册后,即可对含有自定义算子的模型进行执行等操作。

注意

在使用前,请您确认模型的自定义算子名称与注册的算子名称是相同的。

参考文档,请参阅Runtime示例中的 advanced_samples 。

文档内容对你是否有帮助?