Adding a new Op¶
You can create a custom op if it is not supported yet.
To add a custom op, you need to follow these steps:
Implement the Operation¶
The Best way is to refer to the implementation of other operator(e.g. /mace/ops/activation.cc
)
Define the new Op class in mace/ops/my_custom_op.cc
.
- ARM kernels: Kernel about NEON is located at
mace/ops/arm/my_custom_op.cc
- GPU kernels: OpenCL kernel API is defined in
mace/ops/opencl/my_custom_op.h
,- Kernel based on Image is realized in
mace/ops/opencl/image/my_custom_op.cc
, - Kernel based on Buffer is realized in
mace/ops/opencl/buffer/my_custom_op.cc
. - OpenCL kernel file is realized in
mace/ops/opencl/cl/my_custom_op.cl
. - Add the path of opencl kernel file in file
mace/repository/opencl-kernel/opencl_kernel_configure.bzl
- Kernel based on Image is realized in
The structure of Op is like the following code.
#include "mace/core/ops/operator.h"
namespace mace {
namespace ops {
template <RuntimeType D, class T>
class MyCustomOp;
template <>
class MyCustomOp<RuntimeType::RT_CPU, float> : public Operation {
...
}
#ifdef MACE_ENABLE_OPENCL
template<>
class MyCustomOp<RuntimeType::RT_OPENCL, float> : public Operation {
...
};
#endif // MACE_ENABLE_OPENCL
void RegisterMyCustomOp(OpRegistry *op_registry) {
MACE_REGISTER_OP(op_registry, "MyCustomOp", MyCustomOp,
RuntimeType::RT_CPU, float);
MACE_REGISTER_GPU_OP(op_registry, "MyCustomOp", MyCustomOp);
}
} // namespace ops
} // namespace mace
Register the Operation¶
Register the new Op in mace/ops/ops_register.cc
.
#include "mace/ops/ops_register.h"
namespace mace {
namespace ops {
// Keep in lexicographical order
...
extern void RegisterMyCustomOp(OpRegistry *op_registry);
...
} // namespace ops
OpRegistry::OpRegistry() {
// Keep in lexicographical order
...
ops::RegisterMyCustomOp(this);
...
}
} // namespace mace
Add UTs¶
Add operation unit tests in mace/ops/my_custom_op_test.cc
Add benchmark¶
Add operation benchmark in mace/ops/my_custom_op_benchmark.cc
It's strongly recommended to add unit tests and micro benchmarks for your
new Op. If you wish to contribute back, it's required.
Add Op in model converter¶
You need to add this new Op in the model converter.
Document the new Op¶
Finally, add an entry in operator table in the document.