欢迎关注我的公众号 [极智视界]，获取我的更多笔记分享
大家好，我是极智视界。本文介绍了 ncnn 新增 upsample cpu 算子的方法。
ncnn 很强大，提供了丰富的算子实现，但现在网络不断更新，总会涉及到支持不到位的，需要自己动手来写、注册算子，本文就以 upsample cpu 算子新增为例，教你写 ncnn 算子。
下面开始。
编写 upsample 算子实现：
//添加 ncnn::Layer 模板所在的头文件
#include "layer.h"
// 编写算子
class Upsample : public ncnn::Layer{
public:
Upsample(){
one_blob_only = true;
}
virtual int forward(const ncnn::Mat& bottom_blob, ncnn::Mat& top_blob, const ncnn::Option& opt) const
{
int w = bottom_blob.w;
int h = bottom_blob.h;
int channels = bottom_blob.c;
int out_w = w * 2;
int out_h = h * 2;
int out_c = channels;
// 创建输出的 Tensor
top_blob.create(out_w, out_h, out_c, 4u, 1, opt.blob_allocator);
if (top_blob.empty())
return -100;
// openmp 循环展开
#pragma omp parallel for num_threads(opt.num_threads)
for (int p = 0; p < out_c; p++)
{
const float* ptr = bottom_blob.channel(p);
float* outptr = top_blob.channel(p);
for (int i = 0; i < out_h; i++)
{
const float* ptr_h = ptr + i / 2 * w;
float* outptr_h = outptr + i * out_w;
for (int j = 0; j < out_w; j++)
*(outptr_h + j) = *(ptr_h + j / 2);
}
}
return 0;
}
};
使用宏定义生成自定义算子：
# define DEFINE_LAYER_CREATOR(name) \
::ncnn::Layer* name##_layer_creator(void* /*userdata*/) \
{ \
return new name; \
}
DEFINE_LAYER_CREATOR(Upsample)
然后注册到对应的网络中：
ncnn::Net mynet;
// 使用 ncnn::Net::register_custom_layer() 把上述编写的算子注册到 mynet 网络中
mynet.register_custom_layer("Upsample", Upsample_layer_creator);
好了，以上就实现了在 mynet 中新增 upsample cpu 算子，希望我的分享能对你的学习有一点帮助。
【公众号传送】
