caffe源码分析-DataLayer

bear_fish

发布于 2019-02-25 11:53:13

6220

发布于 2019-02-25 11:53:13

DataLayer作为caffe训练时的数据层（以多线程的方式读取数据加速solver的训练过程），继承自BaseDataLayer/BasePrefetchingDataLayer。

template <typename Dtype>
class BaseDataLayer : public Layer<Dtype> {
public:
    explicit BaseDataLayer(const LayerParameter& param);
    // This method may not be overridden except by the BasePrefetchingDataLayer.
    virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top);
    // Data layers should be shared by multiple solvers in parallel
    virtual inline bool ShareInParallel() const { return true; }
    virtual void DataLayerSetUp(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top) {}
    // Data layers have no bottoms, so reshaping is trivial.
    virtual void Reshape(....) {//do nothing}
    virtual void Backward_cpu(....) {//do nothing}

protected:
    TransformationParameter transform_param_;
    shared_ptr<DataTransformer<Dtype> > data_transformer_;
    bool output_labels_;
};

data_transformer_主要是对输入的图片做crop_size，do_mirror，mean等操作.

下面看起函数定义LayerSetUp（参数初始化）：

template <typename Dtype>
void BaseDataLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
                                      const vector<Blob<Dtype>*>& top) {
    if (top.size() == 1) {
        output_labels_ = false;
    } else {
        output_labels_ = true;
    }
    data_transformer_.reset(
            new DataTransformer<Dtype>(transform_param_, this->phase_));
    data_transformer_->InitRand();
    // The subclasses should setup the size of bottom and top
    DataLayerSetUp(bottom, top);
}

BasePrefetchingDataLayer还继承了InternalThread

template <typename Dtype>
class BasePrefetchingDataLayer :
        public BaseDataLayer<Dtype>, public InternalThread {
public:
    explicit BasePrefetchingDataLayer(const LayerParameter& param);

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

    // Prefetches batches 
    static const int PREFETCH_COUNT = 3;

protected:
    virtual void InternalThreadEntry();
    virtual void load_batch(Batch<Dtype>* batch) = 0;

    Batch<Dtype> prefetch_[PREFETCH_COUNT];
    BlockingQueue<Batch<Dtype>*> prefetch_free_;
    BlockingQueue<Batch<Dtype>*> prefetch_full_;

    Blob<Dtype> transformed_data_;
};

template <typename Dtype>
void BasePrefetchingDataLayer<Dtype>::LayerSetUp(
        const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
    BaseDataLayer<Dtype>::LayerSetUp(bottom, top);
    //.........
    DLOG(INFO) << "Initializing prefetch";
    this->data_transformer_->InitRand();
    StartInternalThread();
    DLOG(INFO) << "Prefetch initialized.";
}

其中load_batch函数被DataLayer重载.

template <typename Dtype>
void BasePrefetchingDataLayer<Dtype>::InternalThreadEntry() {
        while (!must_stop()) {
            Batch<Dtype> *batch = prefetch_free_.pop();
            load_batch(batch);

            prefetch_full_.push(batch);
        }
}

DataLayer中使用线程读取Batch（image，label）push到队列中，然后pop出来前向传播：

template <typename Dtype>
void BasePrefetchingDataLayer<Dtype>::Forward_cpu(
        const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {

    Batch<Dtype>* batch = prefetch_full_.pop("Data layer prefetch queue empty");
    // Reshape to loaded data.
    top[0]->ReshapeLike(batch->data_);

    // Copy the data
    caffe_copy(batch->data_.count(), batch->data_.cpu_data(),
               top[0]->mutable_cpu_data());

    if (this->output_labels_) {
        // Reshape to loaded labels.
        top[1]->ReshapeLike(batch->label_);
        // Copy the labels.
        caffe_copy(batch->label_.count(), batch->label_.cpu_data(),
                   top[1]->mutable_cpu_data());
    }

    prefetch_free_.push(batch);
}

下面看DataLayer:

template <typename Dtype>
class DataLayer : public BasePrefetchingDataLayer<Dtype> {
public:
    explicit DataLayer(const LayerParameter& param);
    virtual ~DataLayer();
    virtual void DataLayerSetUp(const vector<Blob<Dtype>*>& bottom,
                                const vector<Blob<Dtype>*>& top);
    // DataLayer uses DataReader instead for sharing for parallelism
    virtual inline bool ShareInParallel() const { return false; }
    virtual inline const char* type() const { return "Data"; }
    virtual inline int ExactNumBottomBlobs() const { return 0; }
    virtual inline int MinTopBlobs() const { return 1; }
    virtual inline int MaxTopBlobs() const { return 2; }

protected:
    virtual void load_batch(Batch<Dtype>* batch);
    DataReader reader_;

};

DataLayerSetUp初始化参数:

template <typename Dtype>
void DataLayer<Dtype>::DataLayerSetUp(const vector<Blob<Dtype>*>& bottom,
                                      const vector<Blob<Dtype>*>& top) {
    const int batch_size = this->layer_param_.data_param().batch_size();
    // Read a data point, and use it to initialize the top blob.
    Datum& datum = *(reader_.full().peek());

    // Use data_transformer to infer the expected blob shape from datum.
    vector<int> top_shape = this->data_transformer_->InferBlobShape(datum);
    this->transformed_data_.Reshape(top_shape);
    // Reshape top[0] and prefetch_data according to the batch_size.
    top_shape[0] = batch_size;
    top[0]->Reshape(top_shape);
    for (int i = 0; i < this->PREFETCH_COUNT; ++i) {
        this->prefetch_[i].data_.Reshape(top_shape);
    }
    
    if (this->output_labels_) {// label
        vector<int> label_shape(1, batch_size);
        top[1]->Reshape(label_shape);
        for (int i = 0; i < this->PREFETCH_COUNT; ++i) {
            this->prefetch_[i].label_.Reshape(label_shape);
        }
    }
}

load_batch在线程中轮训读取数据库(如lmdb)中的数据，并做crop_size，do_mirror，mean等操作，然后通过Forward_cpu传递给下一层.

template<typename Dtype>
void DataLayer<Dtype>::load_batch(Batch<Dtype>* batch) {
    CPUTimer batch_timer;
    batch_timer.Start();
    double read_time = 0;nnndouble trans_time = 0;
    CPUTimer timer;

    const int batch_size = this->layer_param_.data_param().batch_size();
    Datum& datum = *(reader_.full().peek());
    // Use data_transformer to infer the expected blob shape from datum.
    vector<int> top_shape = this->data_transformer_->InferBlobShape(datum);
    this->transformed_data_.Reshape(top_shape);
    // Reshape batch according to the batch_size.
    top_shape[0] = batch_size;
    batch->data_.Reshape(top_shape);

    Dtype* top_data = batch->data_.mutable_cpu_data();
    Dtype* top_label = NULL;

    if (this->output_labels_) {
        top_label = batch->label_.mutable_cpu_data();
    }
    for (int item_id = 0; item_id < batch_size; ++item_id) {
        timer.Start();
        // get a datum
        Datum& datum = *(reader_.full().pop("Waiting for data"));
        read_time += timer.MicroSeconds();
        timer.Start();
        // Apply data transformations (mirror, scale, crop...)
        int offset = batch->data_.offset(item_id);
        this->transformed_data_.set_cpu_data(top_data + offset);
        this->data_transformer_->Transform(datum, &(this->transformed_data_));
        // Copy label.
        if (this->output_labels_) {
            top_label[item_id] = datum.label();
        }
        trans_time += timer.MicroSeconds();

        reader_.free().push(const_cast<Datum*>(&datum));
    }
    timer.Stop();
    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.";
}

使用示例如下：

layer {
  name: "mnist"
  type: "Data"
  top: "data"
  top: "label"
  include {
    phase: TRAIN
  }
  transform_param {
    scale: 0.00390625
  }
  data_param {
    source: "/home/xy/caffe-master/examples/mnist/mnist_train_lmdb"
    batch_size: 64
    backend: LMDB
  }
}

proto定义如下：

message DataParameter {
  enum DB {
    LEVELDB = 0;
    LMDB = 1;
  }
  // Specify the data source.
  optional string source = 1;
  // Specify the batch size.
  optional uint32 batch_size = 4;
  optional uint32 rand_skip = 7 [default = 0];
  optional DB backend = 8 [default = LEVELDB];
  // Prefetch queue (Number of batches to prefetch to host memory.
  optional uint32 prefetch = 10 [default = 4];
}