6.5. 自定义算子开发¶
6.5.1. 简介¶
地平线工具链中已经支持了丰富的算子，在大多数情况下，您的模型应该可以通过前文所述模型转换顺利部署到地平线计算平台上。 已支持的算子情况可以参考 工具链算子支持约束列表 章节。 少部分算子不支持情况下，我们强烈建议您先尝试下替换算子的可能性，这样有利于将地平线计算平台能力充分发挥出来，且开发成本会更低。
自定义算子只提供CPU上算子开发能力，一个完整的自定义算子应用过程包括创建模板、算子实现、算子编译、 含自定义算子模型转换和运行含自定义op模型几个阶段。具体流程如下图所示:
如图所示，定义自定义OP需要有两部分的任务：在模型转换阶段，需要提供自定义OP的python代码；在模拟器/上板运行推理阶段，需要提供自定义OP的C++代码。 要确保这两部分的代码运算的一致性。
6.5.2. 含自定义算子的模型转换¶
6.5.2.1. 模型文件修改¶
在准备好自定义算子实现后，为了将算子应用起来，您需要从原始模型文件和模型转换配置两个方面做出相应调整
(下面分别以 Caffe 模型和 ONNX 模型为例)。
6.5.2.1.1. 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 固定使用 register。
op_register_files 是自定义算子计算的实现文件，请使用相对路径。
完成这些配置后，模型转换的后续步骤与其他一般模型转换过程一致。
6.5.2.1.2. 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中的计算逻辑函数一致
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模型支持多种类型的自定义算子，如您需要在其他框架中支持多种类型的自定义算子请联系地平线技术支持人员。
2.与 Caffe 模型一致，需要在模型转换配置中，使用自定义算子需要在配置文件中加入一个新的自定义op参数组如下：
#...
custom_op:
# 自定义op的校准方式
custom_op_method: register
# 自定义OP的实现文件
op_register_files: sample_custom.py
对于 ONNX 模型，以上参数组中的两个参数都是必须配置的。 custom_op_method 固定使用 register。
op_register_files 是自定义算子计算的实现文件，请使用相对路径。
完成这些配置后，模型转换的后续步骤与其他一般模型转换过程一致。
6.5.2.2. 算子实现¶
在模型转换阶段，需要提供自定义算子的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
@op_implement_register("Cop1")
class Cop1(object):
def __init__(self, ):
pass
def compute(self, x1, x2):
out = x1 + x2 + 1
return out
@op_implement_register("Cop2")
class Cop2(object):
def __init__(self, ):
pass
def compute(self, x1, x2, x3):
out = x1 + x2 + x3 + 1
return out, out
该文件的名字(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模板文件如下:
// custom_identity_add1.h
#ifndef ADVANCED_SAMPLES_CUSTOM_IDENTITY_ADD1_H_
#define ADVANCED_SAMPLES_CUSTOM_IDENTITY_ADD1_H_
#include <string>
#include <vector>
#include "dnn/hb_dnn.h"
#include "dnn/plugin/hb_dnn_layer.h"
#include "dnn/plugin/hb_dnn_ndarray.h"
namespace hobot {
namespace dnn {
Layer *Cop1_layer_creator();
class Cop1 : public Layer {
public:
Cop1() = default;
~Cop1() override = default;
public:
int32_t Init(const Attribute &attributes) override;
int32_t Forward(const std::vector<NDArray *> &bottomBlobs,
std::vector<NDArray *> &topBlobs,
const hbDNNInferCtrlParam *inferCtrlParam) override;
std::string GetType() const override { return "Cop1"; }
uint32_t GetInputCount() const override { return num_args_; }
private:
std::string custom_op_name_;
int32_t num_args_;
};
} // namespace dnn
} // namespace hobot
#endif
// custom_identity_add1.cpp
#include "custom_identity_add1.h"
namespace hobot {
namespace dnn {
Layer *Cop1_layer_creator() { return new Cop1; }
int32_t Cop1::Init(const Attribute &attributes) {
// unused attribute, just demonstrating
attributes.GetAttributeValue(&custom_op_name_, "custom_op_name");
// node's input count
attributes.GetAttributeValue(&num_args_, "num_args");
return 0;
}
int32_t Cop1::Forward(const std::vector<NDArray *> &bottomBlobs,
std::vector<NDArray *> &topBlobs,
const hbDNNInferCtrlParam *inferCtrlParam) {
const NDArray *input0 = bottomBlobs[0];
const NDArray *input1 = bottomBlobs[1];
NDArray *out = topBlobs[0];
const auto *input0_data = input0->Dptr<float>();
const auto *input1_data = input1->Dptr<float>();
auto *out_data = out->Dptr<float>();
uint32_t size = out->Size();
for (uint32_t i = 0U; i < size; i++) {
out_data[i] = input0_data[i] + input1_data[i] + 1;
}
return 0;
}
} // namespace dnn
} // namespace hobot
// custom_identity_add2.h
#ifndef ADVANCED_SAMPLES_CUSTOM_IDENTITY_ADD2_H_
#define ADVANCED_SAMPLES_CUSTOM_IDENTITY_ADD2_H_
#include <string>
#include <vector>
#include "dnn/hb_dnn.h"
#include "dnn/plugin/hb_dnn_layer.h"
#include "dnn/plugin/hb_dnn_ndarray.h"
namespace hobot {
namespace dnn {
Layer *Cop2_layer_creator();
class Cop2 : public Layer {
public:
Cop2() = default;
~Cop2() override = default;
public:
int32_t Init(const Attribute &attributes) override;
int32_t Forward(const std::vector<NDArray *> &bottomBlobs,
std::vector<NDArray *> &topBlobs,
const hbDNNInferCtrlParam *inferCtrlParam) override;
std::string GetType() const override { return "Cop2"; }
uint32_t GetInputCount() const override { return num_args_; }
uint32_t GetOutputCount() const override { return 2U; }
private:
std::string custom_op_name_;
int32_t num_args_;
};
} // namespace dnn
} // namespace hobot
#endif
// custom_identity_add2.cpp
#include "custom_identity_add2.h"
namespace hobot {
namespace dnn {
Layer *Cop2_layer_creator() { return new Cop2; }
int32_t Cop2::Init(const Attribute &attributes) {
// unused attribute, just demonstrating
attributes.GetAttributeValue(&custom_op_name_, "custom_op_name");
// node's input count
attributes.GetAttributeValue(&num_args_, "num_args");
return 0;
}
int32_t Cop2::Forward(const std::vector<NDArray *> &bottomBlobs,
std::vector<NDArray *> &topBlobs,
const hbDNNInferCtrlParam *inferCtrlParam) {
const NDArray *input0 = bottomBlobs[0];
const NDArray *input1 = bottomBlobs[1];
const NDArray *input2 = bottomBlobs[2];
NDArray *out0 = topBlobs[0];
NDArray *out1 = topBlobs[1];
const auto *input0_data = input0->Dptr<float>();
const auto *input1_data = input1->Dptr<float>();
const auto *input2_data = input2->Dptr<float>();
auto *out0_data = out0->Dptr<float>();
auto *out1_data = out1->Dptr<float>();
uint32_t size = out0->Size();
for (uint32_t i = 0U; i < size; i++) {
out0_data[i] = input0_data[i] + input1_data[i] + input2_data[i] + 1;
out1_data[i] = out0_data[i];
}
return 0;
}
} // namespace dnn
} // namespace hobot
注解
该函数名称的前缀(即 Cop1 和 Cop2) 需要与自定义OP的类型（ Kind ）一致, 其传入的参数为：
bottom_blobs：自定义OP节点输入数据。top_blobs：自定义OP节点输出数据。inferCtrlParam：自定义算子初始化阶段的输入参数。
注意
模板中的运算规则均为输出等于所有输入数据累加后再加上数值1，因此后续若需要定义其他行为，则需相应的更改运算规则即可。
6.5.3.2. 自定义算子注册¶
当您完成C++版本模板的修改后，仅需要在示例的CMakeLists.txt中添加对模板文件的包含， 并在应用程序中增加对算子的注册即可，注册请参考以下代码：
#include "custom_identity_add1.h"
#include "custom_identity_add2.h"
hbDNNRegisterLayerCreator("Cop1", hobot::dnn::Cop1_layer_creator);
hbDNNRegisterLayerCreator("Cop2", hobot::dnn::Cop2_layer_creator);
....
当您完成对模板文件的依赖及算子注册后，即可对含有自定义算子的模型进行执行等操作。
注意
在使用前，请您确认模型的自定义算子名称与注册的算子名称是相同的。
参考文档，请参阅Runtime示例中的 advanced_samples 。