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)