caffe添加数据层

0. 添加简单的数据层,生成递增方阵或元素相同的方阵

1. 在src/caffe/proto/caffe.proto添加层参数

  • message LayerParameter中声明新层
  • 在外面定义新层参数,示例如下
...
message LayerParameter {
  optional string name = 1; // the layer name
  optional string type = 2; // the layer type
  repeated string bottom = 3; // the name of each bottom blob
  repeated string top = 4; // the name of each top blob
  ...
  optional ManualDataParameter manual_data_param = 256;
}

message ManualDataParameter {
  optional uint32 size = 1;
  optional string mode = 2; // same increase 
  optional uint32 value = 3;
  optional uint32 step = 4;
}
...

2. 在include/caffe/layers中添加manual_data_layer.hpp

#ifndef MANUAL_DATA_LAYER
#define MANUAL_DATA_LAYER

#include <string>
#include "caffe/blob.hpp"
#include "caffe/layer.hpp"
#include "caffe/proto/caffe.pb.h"

namespace caffe {
    template <typename Dtype>
    class ManualDataLayer : public Layer<Dtype> {
    public:
        explicit ManualDataLayer(const LayerParameter& param)
                : Layer<Dtype>(param) {}
        virtual ~ManualDataLayer(){}
        virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
                                    const vector<Blob<Dtype>*>& top);
        virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
                             const vector<Blob<Dtype>*>& top) {}
        virtual inline const char* type() const { return "ManualData"; }
        virtual inline int ExactNumBottomBlobs() const { return 0; }
        virtual inline int ExactNumTopBlobs() const { return 1; }
    protected:
        virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
                                 const vector<Blob<Dtype>*>& top);
        virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
                                 const vector<Blob<Dtype>*>& top);
        virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
                                  const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {}
        virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
                                  const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {}
        int size_;
        string mode_;
        int value_;
        int step_;                          
    };
}

#endif // !MANUAL_DATA_LAYER

3. 在src/caffe/layers中添加manual_data_layer.cpp

#include "caffe/layers/manual_data_layer.hpp"

namespace caffe {
    template <typename Dtype>
    void ManualDataLayer<Dtype>::LayerSetUp(const std::vector<caffe::Blob<Dtype> *> &bottom,
                    const std::vector<caffe::Blob<Dtype> *> &top) {
        size_ = this->layer_param_.manual_data_param().size();
        top[0]->Reshape(1, 1, size_, size_);
        mode_ = this->layer_param_.manual_data_param().mode();
        value_ = this->layer_param_.manual_data_param().value();
    }

    template <typename Dtype>
    void ManualDataLayer<Dtype>::Forward_cpu(const std::vector<caffe::Blob<Dtype> *> &bottom, 
                    const std::vector<caffe::Blob<Dtype> *> &top) {
        Dtype *data = top[0]->mutable_cpu_data();
        caffe::caffe_set(top[0]->count(), Dtype(0), data);
        if(mode_ == "same") {
            for(int i = 0; i < top[0]->count(); ++i) {
                data[i] = value_;
            }
        } else if(mode_ == "increase") {
            step_ = this->layer_param_.manual_data_param().step();
            for(int i = 0; i < top[0]->count(); ++i) {
                data[i] = value_ + i * step_;
            }
        } else {
            LOG(INFO) << "unknown mode";
            exit(0);
        }
    }

    template <typename Dtype>
    void ManualDataLayer<Dtype>::Forward_gpu(const std::vector<caffe::Blob<Dtype> *> &bottom, 
                    const std::vector<caffe::Blob<Dtype> *> &top) {
        Blob<Dtype> tmp(1, 1, size_, size_);
        Dtype *data = tmp.mutable_cpu_data();
        if(mode_ == "same") {
            for(int i = 0; i < top[0]->count(); ++i) {
                data[i] = value_;
            }
        } else if(mode_ == "increase") {
            step_ = this->layer_param_.manual_data_param().step();
            for(int i = 0; i < top[0]->count(); ++i) {
                data[i] = value_ + i * step_;
            }
        } else {
            LOG(FATAL) << "unknown mode";
        }
        Dtype* gpu_data = top[0]->mutable_gpu_data();
        caffe::caffe_copy(top[0]->count(), tmp.cpu_data(), gpu_data);
    }

    INSTANTIATE_LAYER_GPU_FUNCS(ManualDataLayer);

#ifdef CPU_ONLY
    STUB_GPU_FORWARD(ManualDataLayer, Forward);
#endif

    INSTANTIATE_CLASS(ManualDataLayer);
    REGISTER_LAYER_CLASS(ManualData);
}

4. 使用

layer {
  type: "ManualData"
  name: "data"
  top: "data"
  manual_data_param {
    size: 4
    mode: "increase"
    value: 1
    step: 1
  }
}
发布了50 篇原创文章 · 获赞 31 · 访问量 3万+

猜你喜欢

转载自blog.csdn.net/random_repick/article/details/80601659
今日推荐