专栏算法工具链自定义算子开发

自定义算子开发

颜值即正义2022-03-16
434
2

地平线工具链中已经支持了丰富的算子,在大多数情况下,您的模型应该可以通过使用hb_mapper工具完成转换并顺利部署到地平线芯片上。 少部分算子不支持情况下,我们建议您先尝试下替换算子的可能性,这样有利于将地平线芯片能力充分发挥出来。

自定义算子目前只提供CPU算子开发能力,可自定义onnx算子以及caffe算子。一个完整的自定义算子应用过程包括注册算子、算子实现、含自定义算子模型转换和运行含自定义op模型四个阶段。

1 自定义onnx算子

1.1 将含有自定义算子的pytorch模型导出ONNX

使用torch.onnx.register_custom_op_symbolic注册自定义算子,再导出onnx模型。有以下几处配置参数需要注意:

1. register_custom_op_symbolic函数的第一个参数'::adaptive_avg_pool2d'为pytorch对应操作符名称,若填写错误,则会导致自定义算子注册失败

2. 操作域必须设置为horizon.custom,算子类型为PyOp

3. class_name_s需要与算子实现文件中的类名相对应

4. module_s与算子实现文件名相同,若算子实现文件在当前目录的子目录(custom_op)中,要将相对路径包含进去:"custom_op/sample_custom"

5. 必须指定input_types_i、output_types_i、output_shape_s三个参数

6. 注意指定opset_version为10或11

参考代码:

1.2 算子实现

对应上一节注册自定义算子时配置的算子实现文件(class_name需要保持一致)。

2 自定义caffe算子

2.1 修改prototxt

在原始模型文件中,将自定义算子对应的类型标记为"Custom" ,并设置custom_param。params 是算子的传入参数,指定方式为‘param_name’:param_value, 多个参数之间使用 \n 分隔。

2.2 算子实现

相比于onnx模型,caffe模型的自定义算子实现还需要提供该算子的输出尺寸。

3 含自定义算子的模型转换

在模型转换配置文件中,添加自定义算子相关参数,示例如下:

custom_op_method:固定使用 register
op_register_files:自定义算子计算的实现文件,如果有多份实现,使用 ‘;’ 将各个文件分开即可。

4 含自定义算子的模型推理

想将包含自定算子的.bin模型顺利部署到开发板上,还需要提供自定义算子的C++代码实现。 您可以使用下文提供的模板进行修改:

头文件:

cpp文件:

将以上两个文件放在当前工程目录下之后,编写infer代码时仅需要在加载模型之前增加对算子的注册即可,注册可参考以下代码:

算法工具链
杂谈
评论2
0/1000
  • baek
    Lv.1

    ?

    2022-03-17
    0
    0
  • 青芒-阿白
    Lv.1

    请问J6M支持自定义onnx算子吗

    2025-05-30
    0
    0