地平線工具鏈中已經(jīng)支持了豐富的算子，在大多數(shù)情況下，您的模型應(yīng)該可以通過使用hb_mapper工具完成轉(zhuǎn)換并順利部署到地平線芯片上。 少部分算子不支持情況下，我們建議您先嘗試下替換算子的可能性，這樣有利于將地平線芯片能力充分發(fā)揮出來。

自定義算子目前只提供CPU算子開發(fā)能力，可自定義onnx算子以及caffe算子。一個(gè)完整的自定義算子應(yīng)用過程包括注冊算子、算子實(shí)現(xiàn)、含自定義算子模型轉(zhuǎn)換和運(yùn)行含自定義op模型四個(gè)階段。

1 自定義onnx算子

1.1 將含有自定義算子的pytorch模型導(dǎo)出ONNX

使用torch.onnx.register_custom_op_symbolic注冊自定義算子，再導(dǎo)出onnx模型。有以下幾處配置參數(shù)需要注意：

1. register_custom_op_symbolic函數(shù)的第一個(gè)參數(shù)'::adaptive_avg_pool2d'為pytorch對應(yīng)操作符名稱，若填寫錯(cuò)誤，則會導(dǎo)致自定義算子注冊失敗

2. 操作域必須設(shè)置為horizon.custom，算子類型為PyOp

3. class_name_s需要與算子實(shí)現(xiàn)文件中的類名相對應(yīng)

4. module_s與算子實(shí)現(xiàn)文件名相同，若算子實(shí)現(xiàn)文件在當(dāng)前目錄的子目錄（custom_op）中，要將相對路徑包含進(jìn)去："custom_op/sample_custom"

5. 必須指定input_types_i、output_types_i、output_shape_s三個(gè)參數(shù)

6. 注意指定opset_version為10或11

參考代碼：

import torch

from horizon_nn.horizon_onnx.onnx_pb import TensorProto

from torch.onnx.utils import register_custom_op_symbolic

#prepare your model and input_data

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

register_custom_op_symbolic('::adaptive_avg_pool2d',

horizon_pool,

opset_version=11)

torch.onnx.export(model, input_data, "custom_op.onnx", opset_version=11)

1.2 算子實(shí)現(xiàn)

對應(yīng)上一節(jié)注冊自定義算子時(shí)配置的算子實(shí)現(xiàn)文件（class_name需要保持一致）。

#sample_custom.py

import numpy as np

from horizon_nn.custom import op_implement_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_implement_register("GlobalAveragePool")

class GlobalAveragePool(object):

def __init__(self):

pass

def compute(self, X):

return np.nanmean(X, axis=(2, 3)).reshape(-1, 1024, 1, 1)

2 自定義caffe算子

2.1 修改prototxt

在原始模型文件中，將自定義算子對應(yīng)的類型標(biāo)記為"Custom" ，并設(shè)置custom_param。params 是算子的傳入?yún)?shù)，指定方式為‘param_name’：param_value， 多個(gè)參數(shù)之間使用 \n 分隔。

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"

}

}

2.2 算子實(shí)現(xiàn)

相比于onnx模型，caffe模型的自定義算子實(shí)現(xiàn)還需要提供該算子的輸出尺寸。

#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):

"""Infer the output shapes of the custom operator.

Arguments:

input_shapes: A list of input shapes.

Returns:

Return a list of custom operator's output shapes.

"""

outputs_shape = inputs_shape

return outputs_shape

3 含自定義算子的模型轉(zhuǎn)換

在模型轉(zhuǎn)換配置文件中，添加自定義算子相關(guān)參數(shù)，示例如下：

custom_op_method：固定使用 register；

op_register_files：自定義算子計(jì)算的實(shí)現(xiàn)文件，如果有多份實(shí)現(xiàn)，使用 ‘;’ 將各個(gè)文件分開即可。

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

想將包含自定算子的.bin模型順利部署到開發(fā)板上，還需要提供自定義算子的C++代碼實(shí)現(xiàn)。 您可以使用下文提供的模板進(jìn)行修改：

頭文件：

// custom_identity.h

#ifndef ADVANCED_SAMPLES_CUSTOM_IDENTITY_H_

#define ADVANCED_SAMPLES_CUSTOM_IDENTITY_H_

#include

#include

#include "dnn/hb_dnn.h"

#include "dnn/plugin/hb_dnn_layer.h"

#include "dnn/plugin/hb_dnn_ndarray.h"

namespace hobot {

namespace dnn {

Layer *CustomIdentity_layer_creator();

class CustomIdentity : public Layer {

public:

CustomIdentity() = default;

~CustomIdentity() override = default;

public:

int32_t Init(const Attribute &attributes) override;

int32_t Forward(const std::vector &bottomBlobs,

std::vector &topBlobs,

const hbDNNInferCtrlParam *inferCtrlParam) override;

std::string GetType() const override { return "CustomIdentity"; }

private:

std::string module_;

};

} // namespace dnn

} // namespace hobot

#endif

cpp文件：

// custom_identity.cpp

#include "custom_identity.h"

namespace hobot {

namespace dnn {

Layer *CustomIdentity_layer_creator() { return new CustomIdentity; }

int32_t CustomIdentity::Init(const Attribute &attributes) {

// unused attribute, just demonstrating

attributes.GetAttributeValue(&module_, "module");

return 0;

}

int32_t CustomIdentity::Forward(const std::vector &bottomBlobs,

std::vector &topBlobs,

const hbDNNInferCtrlParam *inferCtrlParam) {

const NDArray *input = bottomBlobs[0];

NDArray *out = topBlobs[0];

const auto *input_data = input->Dptr();

auto *out_data = out->Dptr();

uint32_t size = input->Size();

for (uint32_t i = 0U; i < size; i++) {?

out_data[i] = input_data[i];

}

return 0;

}

} // namespace dnn

} // namespace hobot

將以上兩個(gè)文件放在當(dāng)前工程目錄下之后，編寫infer代碼時(shí)僅需要在加載模型之前增加對算子的注冊即可，注冊可參考以下代碼：

//infer.cpp

#include "custom_identity.h"

// register custom layer

hbDNNRegisterLayerCreator("CustomIdentity",

hobot::dnn::CustomIdentity_layer_creator)

本文轉(zhuǎn)載自地平線開發(fā)者社區(qū)：https://developer.horizon.ai
原作者：顏值即正義
原文鏈接：https://developer.horizon.ai/forumDetail/71036525692881018