caffe 添加DataLayer

Caffe DataLayer 类型层添加说明及数据流说明：

一、 数据流 

Layer->DataLayer->BaseDataLayer->BasePrefectingDataLayer->ImageDimPreFecfchingDataLayer-> 你的类

其中：

Layer 层：（ layer.hpp ）

主要定义了之后层的各种虚函数，包括： LayerSetUp 、 Reshape 、 Forward_cpu\gpu 、 Backward_cpu\gpu 。

其完成的任务为：将 layer_param 中的参数读入内存，并进行 top 、 bottom 参数 size 的检查。具体如下：

  explicit Layer(const LayerParameter& param)

    : layer_param_(param) {

      // The only thing we do is to copy blobs if there are any.

      if (layer_param_.blobs_size() > 0) {

        blobs_.resize(layer_param_.blobs_size());

        for (int i = 0; i < layer_param_.blobs_size(); ++i) {

          blobs_[i].reset(new Blob<Dtype>());

          blobs_[i]->FromProto(layer_param_.blobs(i));

        }

      }

    }

  virtual ~Layer() {}

  /**

   * Called by the parent Layer's SetUp to check that the number of bottom

   * and top Blobs provided as input match the expected numbers specified by

   * the {ExactNum,Min,Max}{Bottom,Top}Blobs() functions.

   */

  virtual void CheckBlobCounts(const vector<Blob<Dtype>*>& bottom,

                               const vector<Blob<Dtype>*>& top) {

    if (ExactNumBottomBlobs() >= 0) {

      CHECK_EQ(ExactNumBottomBlobs(), bottom.size())

          << type_name() << " Layer takes " << ExactNumBottomBlobs()

          << " bottom blob(s) as input.";

    }

    if (MinBottomBlobs() >= 0) {

      CHECK_LE(MinBottomBlobs(), bottom.size())

          << type_name() << " Layer takes at least " << MinBottomBlobs()

          << " bottom blob(s) as input.";

    }

    if (MaxBottomBlobs() >= 0) {

      CHECK_GE(MaxBottomBlobs(), bottom.size())

          << type_name() << " Layer takes at most " << MaxBottomBlobs()

          << " bottom blob(s) as input.";

    }

    if (ExactNumTopBlobs() >= 0) {

      //cout top.size() --add by longmao

      //std::cout<<"ExactNumTopBlobs(): "<<ExactNumTopBlobs()<<"   top.size(): "<<top.size()<<std::endl;

      CHECK_EQ(ExactNumTopBlobs(), top.size())

          << type_name() << " Layer produces " << ExactNumTopBlobs()

          << " top blob(s) as output.";

    }

    if (MinTopBlobs() >= 0) {

      CHECK_LE(MinTopBlobs(), top.size())

          << type_name() << " Layer produces at least " << MinTopBlobs()

          << " top blob(s) as output.";

    }

    if (MaxTopBlobs() >= 0) {

      CHECK_GE(MaxTopBlobs(), top.size())

          << type_name() << " Layer produces at most " << MaxTopBlobs()

          << " top blob(s) as output.";

    }

    if (EqualNumBottomTopBlobs()) {

      CHECK_EQ(bottom.size(), top.size())

          << type_name() << " Layer produces one top blob as output for each "

          << "bottom blob input.";

    }

  }

BaseLayer(base_data_layers.cpp 、 data_layers.hpp)

主要实现：

         BaseDataLayer ：构造函数，推测将 param 向下传递

         LayerSetp ：使得上层可以调用底层 DataLayerSetUp

BasePrefetchingDataLayer(base_data_layers.cpp 、 data_layers.hpp)

主要实现：

LayerSetUp ： 调用 CreatePrefetchThread 开启数据读写线程

CreatePrefetchThread 、 JoinPrefectchThread ：线程操作函数

Forward_cpu ：调用 caffe_copy 将数据放入指针地址中

         

         ImageDimPrefetchingDataLayer （ base_data_layers.cpp 、 data_layers.hpp ）

          主要实现：

                   LayerSetUp ： 申请内存，并建立指针地址相关信息

Forward_cpu ：调用 caffe_copy 将数据放入指针地址中

二、增加 layer

          由于本次增加 layer 功能为将 image 及 float 同时作为 label 向下传入，因此继承 ImageDimPrefetchingDataLayer 层。

          主要实现：

ImageDimPrefetchingDataLayer ：从 param 中读取 label 数据，存入 lines_ 中，并开辟内存

                   InternalTreadEntry ：交由上层调用，将 lines_ 中的数据通过 set_cpu_data 放入内存，交由上层 forward_cpu 进行拷贝。

          代码如下：

data_layers.hpp

template <typename Dtype>

class ImageSegAtmDataLayer : public ImageDimPrefetchingDataLayer<Dtype> {

 public:

  explicit ImageSegAtmDataLayer(const LayerParameter& param)

    : ImageDimPrefetchingDataLayer<Dtype>(param) {}

  virtual ~ImageSegAtmDataLayer();

  virtual void DataLayerSetUp(const vector<Blob<Dtype>*>& bottom,

      const vector<Blob<Dtype>*>& top);

  virtual inline LayerParameter_LayerType type() const {

    return LayerParameter_LayerType_IMAGE_DATA;

  }

  virtual inline int ExactNumBottomBlobs() const { return 0; }

  // related to the output top numbers for init and set up memory

  virtual inline int ExactNumTopBlobs() const { return 5; }

  virtual inline bool AutoTopBlobs() const { return true; }

 

protected:

  virtual void ShuffleImages();

  virtual void InternalThreadEntry();

 

protected:

  Blob<Dtype> transformed_label_;

  // add transformed_atm_*_ for ImageSegAtmDataLayer  --add by longmao

  Blob<Dtype> transformed_atm_r_;

  Blob<Dtype> transformed_atm_g_;

  Blob<Dtype> transformed_atm_b_;

  shared_ptr<Caffe::RNG> prefetch_rng_;

  typedef struct SegItems {

    std::string imgfn;

    std::string segfn;

    float r,g,b;

  } SEGITEMS;

  vector<SEGITEMS> lines_;

  int lines_id_;

};

image_seg_atm_data_layer.cpp

/*

 notice:

 this code is based on the following implementation.

 https://bitbucket.org/deeplab/deeplab-public/

 */

#include <fstream>  // NOLINT(readability/streams)

#include <iostream>  // NOLINT(readability/streams)

#include <string>

#include <utility>

#include <vector>

#include <algorithm>

#include <opencv2/core/core.hpp>

#include <opencv2/highgui/highgui.hpp>

#include <opencv2/highgui/highgui_c.h>

#include <opencv2/imgproc/imgproc.hpp>

#include "caffe/data_layers.hpp"

#include "caffe/layer.hpp"

#include "caffe/util/benchmark.hpp"

#include "caffe/util/io.hpp"

#include "caffe/util/math_functions.hpp"

#include "caffe/util/rng.hpp"

namespace caffe {

template <typename Dtype>

ImageSegAtmDataLayer<Dtype>::~ImageSegAtmDataLayer<Dtype>() {

  this->JoinPrefetchThread();

}

// lines_               --a vector to stor data from file

// param                --the address of file while append from net.cpp

// prefetch_data_       --set memory for data

// transformed_label_   --processing data such as copy data to cpu_data

// reshape              --set up blobe`s dimesion

// 注：由于本层没有实现forward_cpu（将数据copy到cpu中）方法 具体到其父类ImageDimPrefetchingDataLayer 内实现

// read data from txt and set the memory --add by longmao

template <typename Dtype>

void ImageSegAtmDataLayer<Dtype>::DataLayerSetUp(const vector<Blob<Dtype>*>& bottom,

      const vector<Blob<Dtype>*>& top) {

  LOG(INFO) << "DataLayerSetUp by longamo";

  const int new_height = this->layer_param_.image_data_param().new_height();

  const int new_width  = this->layer_param_.image_data_param().new_width();

  const bool is_color  = this->layer_param_.image_data_param().is_color();

  const int label_type = this->layer_param_.image_data_param().label_type();

  string root_folder = this->layer_param_.image_data_param().root_folder();

  TransformationParameter transform_param = this->layer_param_.transform_param();

  CHECK(transform_param.has_mean_file() == false) <<

         "ImageSegAtmDataLayer does not support mean file";

  CHECK((new_height == 0 && new_width == 0) ||

      (new_height > 0 && new_width > 0)) << "Current implementation requires "

      "new_height and new_width to be set at the same time.";

  // Read the file with filenames and labels

  const string& source = this->layer_param_.image_data_param().source();

  LOG(INFO) << "Opening file " << source;

  std::ifstream infile(source.c_str());

  string linestr;

  //read data --add by longmao

  while (std::getline(infile, linestr)) {

    std::istringstream iss(linestr);

    string imgfn;

    iss >> imgfn;

    string segfn = "";

    if (label_type != ImageDataParameter_LabelType_NONE) {

      iss >> segfn;

    }

    // add struct to read data --add by lyn

    SEGITEMS item;

    item.imgfn = imgfn;

    item.segfn = segfn;

    float r,g,b;

    iss >> r >> g >> b;

    item.r = r;

    item.g = g;

    item.b = b;

    lines_.push_back(item);

  }

  // get memory --add by

  if (this->layer_param_.image_data_param().shuffle()) {

    // randomly shuffle data

    LOG(INFO) << "Shuffling data";

    const unsigned int prefetch_rng_seed = caffe_rng_rand();

    prefetch_rng_.reset(new Caffe::RNG(prefetch_rng_seed));

    ShuffleImages();

  }

  LOG(INFO) << "A total of " << lines_.size() << " images.";

  lines_id_ = 0;

  // Check if we would need to randomly skip a few data points

  if (this->layer_param_.image_data_param().rand_skip()) {

    unsigned int skip = caffe_rng_rand() %

        this->layer_param_.image_data_param().rand_skip();

    LOG(INFO) << "Skipping first " << skip << " data points.";

    CHECK_GT(lines_.size(), skip) << "Not enough points to skip";

    lines_id_ = skip;

  }

  // Read an image, and use it to initialize the top blob.

  cv::Mat cv_img = ReadImageToCVMat(root_folder + lines_[lines_id_].imgfn,

                                    new_height, new_width, is_color);

  const int channels = cv_img.channels();

  const int height = cv_img.rows;

  const int width = cv_img.cols;

  // image

  const int crop_size = this->layer_param_.transform_param().crop_size();

  const int batch_size = this->layer_param_.image_data_param().batch_size();

  if (crop_size > 0) {

    top[0]->Reshape(batch_size, channels, crop_size, crop_size);

    this->prefetch_data_.Reshape(batch_size, channels, crop_size, crop_size);

    this->transformed_data_.Reshape(1, channels, crop_size, crop_size);

    // label

    top[1]->Reshape(batch_size, 1, crop_size, crop_size);

    this->prefetch_label_.Reshape(batch_size, 1, crop_size, crop_size);

    this->transformed_label_.Reshape(1, 1, crop_size, crop_size);

  } else {

    top[0]->Reshape(batch_size, channels, height, width);

    this->prefetch_data_.Reshape(batch_size, channels, height, width);

    this->transformed_data_.Reshape(1, channels, height, width);

    // label

    top[1]->Reshape(batch_size, 1, height, width);

    this->prefetch_label_.Reshape(batch_size, 1, height, width);

    this->transformed_label_.Reshape(1, 1, height, width);

  }

  // data

  LOG(INFO) << "output data size: " << top[0]->num() << ","

      << top[0]->channels() << "," << top[0]->height() << ","

      << top[0]->width();

  // label

  LOG(INFO) << "output label size: " << top[1]->num() << ","

      << top[1]->channels() << "," << top[1]->height() << ","

      << top[1]->width();

  // atm_r --add by longmao

  top[2]->Reshape(batch_size, 1, 1, 1);

  this->prefetch_atm_r_.Reshape(batch_size, 1, 1, 1);

  this->transformed_atm_r_.Reshape(1,1,1,1);

  LOG(INFO) << "output data_dim size: " << top[2]->num() << ","

      << top[2]->channels() << "," << top[2]->height() << ","

      << top[2]->width();

  // atm_g --add by longmao

  top[3]->Reshape(batch_size, 1, 1, 1);

  this->prefetch_atm_g_.Reshape(batch_size, 1, 1, 1);

  this->transformed_atm_g_.Reshape(1,1,1,1);

  LOG(INFO) << "output Ar size: " << top[3]->num() << ","

      << top[3]->channels() << "," << top[3]->height() << ","

      << top[3]->width() ;

  // atm_b --add by longmao

  top[4]->Reshape(batch_size, 1, 1, 1);

  this->prefetch_atm_b_.Reshape(batch_size, 1, 1, 1);

  this->transformed_atm_b_.Reshape(1,1,1,1);

  LOG(INFO) << "output Ar size: " << top[4]->num() << ","

      << top[4]->channels() << "," << top[4]->height() << ","

      << top[4]->width() ;

}

template <typename Dtype>

void ImageSegAtmDataLayer<Dtype>::ShuffleImages() {

  caffe::rng_t* prefetch_rng =

      static_cast<caffe::rng_t*>(prefetch_rng_->generator());

  shuffle(lines_.begin(), lines_.end(), prefetch_rng);

}

// This function is used to create a thread that prefetches and transform the data.

template <typename Dtype>

void ImageSegAtmDataLayer<Dtype>::InternalThreadEntry() {

  //LOG(INFO) << "InternalThreadEntry";

  CPUTimer batch_timer;

  batch_timer.Start();

  double read_time = 0;

  double trans_time = 0;

  CPUTimer timer;

  CHECK(this->prefetch_data_.count());

  CHECK(this->transformed_data_.count());

  //CHECK(this->transformed_atm_r_.count());

  //

  Dtype* top_data     = this->prefetch_data_.mutable_cpu_data();

  Dtype* top_label    = this->prefetch_label_.mutable_cpu_data();

  // get head pointer from mutable_cpu_data about atm_* --add by longmao

  Dtype* top_atm_r       = this->prefetch_atm_r_.mutable_cpu_data();

  Dtype* top_atm_g       = this->prefetch_atm_g_.mutable_cpu_data();

  Dtype* top_atm_b       = this->prefetch_atm_b_.mutable_cpu_data();

  ImageDataParameter image_data_param  = this->layer_param_.image_data_param();

  const int batch_size = image_data_param.batch_size();

  const int new_height = image_data_param.new_height();

  const int new_width  = image_data_param.new_width();

  const int label_type = this->layer_param_.image_data_param().label_type();

  const int ignore_label = image_data_param.ignore_label();

  const bool is_color  = image_data_param.is_color();

  string root_folder   = image_data_param.root_folder();

  const int lines_size = lines_.size();

  // offset of top_atm_* --add by longmao

  int top_atm_r_offset;

  int top_atm_g_offset;

  int top_atm_b_offset;

  for (int item_id = 0; item_id < batch_size; ++item_id) {

    // get offset of top_atm_* --add by longamo

    top_atm_r_offset = this->prefetch_atm_r_.offset(item_id);

    top_atm_g_offset = this->prefetch_atm_g_.offset(item_id);

    top_atm_b_offset = this->prefetch_atm_b_.offset(item_id);

    // read data from lines_ to top_atm_* --add by longmao

    top_atm_r[top_atm_r_offset] = lines_[lines_id_].r;

    top_atm_g[top_atm_g_offset] = lines_[lines_id_].g;

    top_atm_b[top_atm_b_offset] = lines_[lines_id_].b;

    std::vector<cv::Mat> cv_img_seg;

    cv::Mat cv_img, cv_seg;

    // get a blob

    timer.Start();

    CHECK_GT(lines_size, lines_id_);

    int img_row, img_col;

    cv_img = ReadImageToCVMat(root_folder + lines_[lines_id_].imgfn,

    0, 0, is_color, &img_row, &img_col);

    if (!cv_img.data) {

      DLOG(INFO) << "Fail to load img: " << root_folder + lines_[lines_id_].imgfn;

    }

    if (label_type == ImageDataParameter_LabelType_PIXEL) {

      cv_seg = ReadImageToCVMatNearest(root_folder + lines_[lines_id_].segfn,

              0, 0, false);

      if (!cv_seg.data) {

        DLOG(INFO) << "Fail to load seg: " << root_folder + lines_[lines_id_].segfn;

      }

    }

    else if (label_type == ImageDataParameter_LabelType_IMAGE) {

      const int label = atoi(lines_[lines_id_].segfn.c_str());

      cv::Mat seg(cv_img.rows, cv_img.cols,

      CV_8UC1, cv::Scalar(label));

      cv_seg = seg;

    }

    else {

      cv::Mat seg(cv_img.rows, cv_img.cols,

      CV_8UC1, cv::Scalar(ignore_label));

      cv_seg = seg;

    }

    cv_img_seg.push_back(cv_img);

    cv_img_seg.push_back(cv_seg);

    read_time += timer.MicroSeconds();

    timer.Start();

    // Apply transformations (mirror, crop...) to the image

    int offset;

    offset = this->prefetch_data_.offset(item_id);

    this->transformed_data_.set_cpu_data(top_data + offset);

    offset = this->prefetch_label_.offset(item_id);

    this->transformed_label_.set_cpu_data(top_label + offset);

    // set cpu_data in transformed_atm_*_  --add by lognmao

    // 推测此处时将数据传输至上层类进行处理

    offset = this->prefetch_atm_r_.offset(item_id);

    this->transformed_atm_r_.set_cpu_data(top_atm_r + offset);

    offset = this->prefetch_atm_g_.offset(item_id);

    this->transformed_atm_g_.set_cpu_data(top_atm_g + offset);

    offset = this->prefetch_atm_b_.offset(item_id);

    this->transformed_atm_b_.set_cpu_data(top_atm_b + offset);

    //std::cout<<"top_atm_b + offset: "<<*(top_atm_b + offset)<<std::endl;

    // transforsmer data and label with TransformImgAndSeg --add by longmao

    this->data_transformer_.TransformImgAndSeg(cv_img_seg,

   &(this->transformed_data_), &(this->transformed_label_),

   ignore_label);

    trans_time += timer.MicroSeconds();

    // go to the next std::vector<int>::iterator iter;

    lines_id_++;

    if (lines_id_ >= lines_size) {

      // We have reached the end. Restart from the first.

      DLOG(INFO) << "Restarting data prefetching from start.";

      lines_id_ = 0;

      if (this->layer_param_.image_data_param().shuffle()) {

  ShuffleImages();

      }

    }

  }

  batch_timer.Stop();

  DLOG(INFO) << "Prefetch batch: " << batch_timer.MilliSeconds() << " ms.";

  DLOG(INFO) << "     Read time: " << read_time / 1000 << " ms.";

  DLOG(INFO) << "Transform time: " << trans_time / 1000 << " ms.";

}

INSTANTIATE_CLASS(ImageSegAtmDataLayer);

REGISTER_LAYER_CLASS(IMAGE_SEG_ATM_DATA, ImageSegAtmDataLayer);

}  // namespace caffe

由于没有实现 Forward_cpu 因此修改上层类：

data_layers.hpp

template <typename Dtype>

class ImageDimPrefetchingDataLayer : public BasePrefetchingDataLayer<Dtype> {

/*

 notice:

 this code is based on the following implementation.

 https://bitbucket.org/deeplab/deeplab-public/

 */

 public:

  explicit ImageDimPrefetchingDataLayer(const LayerParameter& param)

      : BasePrefetchingDataLayer<Dtype>(param) {}

  virtual ~ImageDimPrefetchingDataLayer() {}

  // LayerSetUp: implements common data layer setup functionality, and calls

  // DataLayerSetUp to do special data layer setup for individual layer types.

  // This method may not be overridden.

  void LayerSetUp(const vector<Blob<Dtype>*>& bottom,

      const vector<Blob<Dtype>*>& top);

  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);

  // The thread's function

  virtual void InternalThreadEntry() {}

 protected:

  Blob<Dtype> prefetch_data_dim_;

  bool output_data_dim_;

  // add prefetch_atm_*_ which used in image_seg_atm_data_layer  --add by longmao;

  bool output_atm_;

  Blob<Dtype> prefetch_atm_r_;

  Blob<Dtype> prefetch_atm_g_;

  Blob<Dtype> prefetch_atm_b_;

};

同时修改： data_layers.cpp 及 data_layers.cu 的 LaterSetUp 函数和 Forward_cpu 函数