Caffe 源码阅读笔记 [基本模块] Syncedmem & Blob

syncedmem

• syncedmem管理一段大小为size的内存。这段内存可以从GPU或者主机内存分配，syncedmem负责GPU和主机内存之间的同步。如果数据是在GPU里而要从主机内存取出，syncedmem会先把数据从GPU内存memcpy到主机内存返回（标记为SYNCED); 相反也是一样。
• 分配/释放内存的方式
  
  • CPU: ptr = malloc(size) / free(ptr)
  • GPU: cudaMallocHost(ptr, size) / cudaFreeHost(ptr)

成员变量

  void* cpu_ptr_, gpu_ptr_; //数据指针
  size_t size_; // SyncedMemory(size_t size)
  // 状态 {
  // UNINITIALIZED = 初始值
  // HEAD_AT_CPU = 最新数据在主机内存里
  // HEAD_AT_GPU = 最新数据在GPU内存里
  // SYNCED = 表示主机内存数据和GPU内存的数据是一致的
  // }
  SyncedHead head_; 
  bool cpu_malloc_use_cuda_;
  bool own_gpu_data_, own_cpu_data_; // true则析构时将内存释放
  int gpu_device_; // GPU device id

CPU 相关方法

// 获得主机内存数据
inline void SyncedMemory::cpu_data() {
  switch (head_) {
  case UNINITIALIZED:
    分配内存，head_=HEAD_AT_CPU
    break;
  case HEAD_AT_GPU: //数据已经在GPU内存里
#ifndef CPU_ONLY
    // 从GPU内存拷贝回主机内存
    caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_);
    head_ = SYNCED;
#end
  return (const void*)cpu_ptr_;
}
// 设置主机内存数据
void SyncedMemory::set_cpu_data(void* data) {
  cpu_ptr_ = data;
  head_ = HEAD_AT_CPU;
  own_cpu_data_ = false;
}

GPU 相关方法

//获得GPU内存数据
inline void SyncedMemory::gpu_data() {
#ifndef CPU_ONLY
  switch (head_) {
  case UNINITIALIZED:
    分配内存, head_ = HEAD_AT_GPU;
    break;
  case HEAD_AT_CPU:
    // 从Host内存拷贝到GPU内存
    caffe_gpu_memcpy(size_, cpu_ptr_, gpu_ptr_);
    head_ = SYNCED;
    break;
#end
}
// 设置GPU内存
void SyncedMemory::set_gpu_data(void* data) {
  gpu_ptr_ = data;
  head_ = HEAD_AT_GPU;
  own_gpu_data_ = false;
}
// 用于在Data Layer异步从内存拷贝到GPU
void SyncedMemory::async_gpu_push(const cudaStream_t& stream) {
  CHECK(head_ == HEAD_AT_CPU); // 确认数据已经在主机内存里
  // 异步拷贝到GPU内存
  CUDA_CHECK(cudaMemcpyAsync(gpu_ptr_, cpu_ptr_, size_, put, stream));
  head_ = SYNCED; // 设置同步完毕
}

Blob

• Blob是Caffe的Layer之间传输的内部数据单元
• 它存储矩阵data_和 diff_， 由shape_来表示矩阵的各个维度。比如number张图片训练数据可以表示成 number * channels * height * width 的4维矩阵，shape_ = {number, channels, height, width}，具体实现上data_和diff_分别是一个大小为number * channels * height * width的一维数组(SyncedMem类型）。

Blob 成员变量

  shared_ptr<SyncedMemory> data_; // 数据
  shared_ptr<SyncedMemory> diff_; // diff
  vector<int> shape_; // 维度
  int count_; // shape_[0]*shape_[1]*...*shape_[n]
  int capacity_; // data_和diff_的capacity，只有当count_>capacity_的时候，才重新分配内存

Blob基本方法

// 更改Blob的维度shape_，如果已分配的内存无法存储shape表示的数组，重新分配内存
void Reshape(const vector<int>& shape); 
// 返回data_->xpu_data(), data_->set_xpu_data()
void cpu_data(), gpu_data(), set_cpu_data(), set_gpu_data() // cpu_diff(), ...类似
// 和别的blob共享data_和diff_，避免拷贝
void ShareData(const Blob& other); // ShareDiff
// 从ProtoBuf读取Blob和写到ProtoBuf里, toProto可设置是否写diff数据
void FromProto/toProto();

Blob矩阵操作

// 更新主机内存或者GPU内存里的数据：data_ := -1*diff_ + data_
void Update() {
    // 由data_.head() == HEAD_AT_CPU 判断数据在哪里, 下面函数类似
    caffe_axpy<Dtype>(count_, Dtype(-1),
        static_cast<const Dtype*>(diff_->cpu_data()),
        static_cast<Dtype*>(data_->mutable_cpu_data()));
    OR
    caffe_gpu_axpy<Dtype>(count_, Dtype(-1),
        static_cast<const Dtype*>(diff_->gpu_data()),
        static_cast<Dtype*>(data_->mutable_gpu_data()));
}
// 计算sum(data_)，data_内所有元素之和
Dtype asum_data() // 调用caffe_cpu_asum 或 caffe_gpu_asum
Dtype asum_diff()
// 计算sum(data_.*data_)，元素平方总和
Dtype sumsq_data() // 调用caffe_cpu_dot(count_, data, data)或caffe_gpu_dot(count_, data, data, &sumsq)
Dtype sumsq_diff()
// data_ := scale_factor * data_
void scale_data(Dtype scale_factor) // 调用caffe_scal 或 caffe_gpu_scal
void scale_diff(Dtype scale_factor)
// 根据index访问data_和diff_矩阵元素
Dtype data_at(const vector<int>& index)
Dtype diff_at(const vector<int>& index)
// 计算一个矩阵切片的元素个数，返回所有shape_[i]的乘积(start_axis<=i<end_axis)
int count(int start_axis, int end_axis)