模型推理ncnn 新增 upsample cpu 算子方法

Posted 2022-03-16 极智视界

  欢迎关注我的公众号 [极智视界]，获取我的更多笔记分享

  大家好，我是极智视界。本文介绍了 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 算子，希望我的分享能对你的学习有一点帮助。

 【公众号传送】

《【模型推理】ncnn 新增 upsample cpu 算子方法》

扫描下方二维码即可关注我的微信公众号【极智视界】，获取更多AI经验分享，让我们用极致+极客的心态来迎接AI ！