如果遇到框架不支持的Op兼蕊，用戶可以自行創(chuàng)建初厚。MACE提供了非常友好的創(chuàng)建自定義Op的操作，只需按照以下步驟執(zhí)行孙技。
1. 定義Op類
創(chuàng)建新文件mace/ops/my_custom_op.h产禾，并在文件中定義新的類，例如MyCustomOp绪杏，代碼如下
#ifndef MACE_OPS_MY_CUSTOM_OP_H_
#define MACE_OPS_MY_CUSTOM_OP_H_
#include "mace/core/operator.h"
#include "mace/kernels/my_custom_op.h"
namespace mace {
namespace ops {
template <DeviceType D, typename T>
class MyCustomOp : public Operator<D, T> {
public:
MyCustomOp(const OperatorDef &op_def, Workspace *ws)
: Operator<D, T>(op_def, ws),
functor_() {}
bool Run(StatsFuture *future) override {
const Tensor *input = this->Input(INPUT);
Tensor *output = this->Output(OUTPUT);
functor_(input, output, future);
return true;
}
protected:
OP_INPUT_TAGS(INPUT);
OP_OUTPUT_TAGS(OUTPUT);
private:
kernels::MyCustomOpFunctor<D, T> functor_;
};
} // namespace ops
} // namespace mace
#endif // MACE_OPS_MY_CUSTOM_OP_H_
2. 注冊(cè)O(shè)p類
創(chuàng)建新文件 mace/ops/my_custom_op.cc下愈，并在其中注冊(cè)新建的類纽绍，代碼如下蕾久。從代碼里可以看出，注冊(cè)代碼提供了CPU和GPU的兩種實(shí)現(xiàn)拌夏，其中GPU版本包含float和half的兩種數(shù)據(jù)支持僧著。
除此之外，還需要在mace/core/operator.cc文件的namespace ops下按模板添加注冊(cè)代碼extern void Register_My_Custom_Op(OperatorRegistry *op_registry);障簿，同時(shí)在OperatorRegistry()構(gòu)造函數(shù)中添加ops::Register_My_Custom_Op(this);
#include "mace/ops/my_custom_op.h"
namespace mace {
namespace ops {
void Register_My_Custom_Op(OperatorRegistry *op_registry) {
REGISTER_OPERATOR(op_registry, OpKeyBuilder("my_custom_op")
.Device(DeviceType::CPU)
.TypeConstraint<float>("T")
.Build(),
Custom_Op<DeviceType::CPU, float>);
REGISTER_OPERATOR(op_registry, OpKeyBuilder("my_custom_op")
.Device(DeviceType::OPENCL)
.TypeConstraint<float>("T")
.Build(),
Custom_Op<DeviceType::OPENCL, float>);
REGISTER_OPERATOR(op_registry, OpKeyBuilder("my_custom_op")
.Device(DeviceType::OPENCL)
.TypeConstraint<half>("T")
.Build(),
Custom_Op<DeviceType::OPENCL, half>);
}
} // namespace ops
} // namespace mace
3. 實(shí)現(xiàn)Op的核心代碼
注冊(cè)完成Op后盹愚，需要實(shí)現(xiàn)其核心代碼。創(chuàng)建文件mace/kernels/my_custom_op.h站故，這個(gè)文件文件實(shí)現(xiàn)的是CPU版本的代碼皆怕。也可以選擇實(shí)現(xiàn)OpenCL版本代碼，這時(shí)西篓，需要?jiǎng)?chuàng)建兩個(gè)文件：mace/kernels/opencl/my_custom_op_opencl.cc 和 mace/kernels/opencl/cl/ my_custom_op.cl愈腾。也可以對(duì)CPU版本的實(shí)現(xiàn)進(jìn)行NEON指令集優(yōu)化。
