caffe源码阅读,1-数据流Blob

Blob是Caffe中层之间数据流通的单位,各个layer之间的数据通过Blob传递。在看Blob源码之前,先看一下CPU和GPU内存之间的数据同步类SyncedMemory;使用GPU运算时,数据要在GPU显存中,但是一开始数据是通过CPU读到内存,通过类SyncedMemory来实现显存和内存之间的数据的同步。

SyncedMemory

先看一下成员变量

  //数据在cpu或gpu,指向数据的指针
  void* cpu_ptr_;
  void* gpu_ptr_;
  size_t size_;//数据大小
  SyncedHead head_;//数据状态,有四种:UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED
  bool own_cpu_data_;
  bool cpu_malloc_use_cuda_;//是否使用cuda标记
  bool own_gpu_data_;
  int gpu_device_;

两个指针分别指向在内存和显存的数据,size_记录数据大小,head_是枚举变量,记录数据状态。gpu_device_指出使用哪块显卡。

成员函数根据名字能看出大概意思,其中

void async_gpu_push(const cudaStream_t& stream);

是异步同步数据到GPU,这里的"异步“是指把数据同步到GPU,在同步未完成时就返回,不需要等待完成同步。

CaffeMallocHost/CaffeFreeHost

这是一个功能和malloc/free相同的分配/释放内存/显存的函数。如果使用了GPU,则在在GPU上分配和释放,否则在内存上分配和释放。

Blob

Blob类的成员变量很少

protected:
  shared_ptr<SyncedMemory> data_;//存放数据
  shared_ptr<SyncedMemory> diff_;//存放梯度
  shared_ptr<SyncedMemory> shape_data_;//Blob形状,N K H W
  vector<int> shape_;//保存 N K H W
  int count_;//元素个数
  int capacity_;//当前元素个数

Blob存储着图像数据,以及偏差。图像数据大小由channel、height、width判断,一个Blob可能存储多幅图像,所以多了一个num。即Blob大小有Num,K(channel),Height,Weight决定。

Blob成员函数很多:

Reshape函数用来调整Blob形状,最终调用的函数如下

 template <typename Dtype>
void Blob<Dtype>::Reshape(const vector<int>& shape) {
  CHECK_LE(shape.size(), kMaxBlobAxes);//维数不能超过kMaxBlobAxes
  count_ = 1;//赋值为1,为了相乘
  shape_.resize(shape.size());
  if (!shape_data_ || shape_data_->size() < shape.size() * sizeof(int)) {
    shape_data_.reset(new SyncedMemory(shape.size() * sizeof(int)));
  }
  int* shape_data = static_cast<int*>(shape_data_->mutable_cpu_data());
  for (int i = 0; i < shape.size(); ++i) {
    CHECK_GE(shape[i], 0);
    CHECK_LE(shape[i], INT_MAX / count_) << "blob size exceeds INT_MAX";
    count_ *= shape[i];//记录数据大小
    shape_[i] = shape[i];
    shape_data[i] = shape[i];
  }
  if (count_ > capacity_) {//capactity不小于count
    capacity_ = count_;
    data_.reset(new SyncedMemory(capacity_ * sizeof(Dtype)));
    diff_.reset(new SyncedMemory(capacity_ * sizeof(Dtype)));
  }
}

可以看出,如果Reshape时,如果大小不够时,会重新分配内存/显存,释放原有内存/显存。

Update()函数用来更新数据,根据数据所在位置进行更新

 template <typename Dtype>
void Blob<Dtype>::Update() {
  // We will perform update based on where the data is located.
  switch (data_->head()) {
  case SyncedMemory::HEAD_AT_CPU:
    // perform computation on CPU
    //data_ = data_ - diff_
    caffe_axpy<Dtype>(count_, Dtype(-1),
        static_cast<const Dtype*>(diff_->cpu_data()),
        static_cast<Dtype*>(data_->mutable_cpu_data()));
    break;
  case SyncedMemory::HEAD_AT_GPU:
  case SyncedMemory::SYNCED:
#ifndef CPU_ONLY
    // perform computation on GPU
    //data_ = data_ - diff_
    caffe_gpu_axpy<Dtype>(count_, Dtype(-1),
        static_cast<const Dtype*>(diff_->gpu_data()),
        static_cast<Dtype*>(data_->mutable_gpu_data()));
#else
    NO_GPU;
#endif
    break;
  default:
    LOG(FATAL) << "Syncedmem not initialized.";
  }
}

Update()函数,实际上进行的运算时data_ = data_ - diff_。

计算范数的函数是特化实现的:

函数asum_data()asum_diff()是计算data_或diff_的L1范数。

函数sumsq_data()sumsq_diff()是计算data_或diff_的L1范数。

其他函数根据名字都可以大概理解了。