Caffe源碼(三)—— Blob

Caffe 源碼 —— blob.hpp/cpp

syncedmem.hpp/cpp

在介紹 Blob 之前得先說一說 syncedmem.hpp/cpp,syncedmem 文件中
定義了用於數據 CPU 和 GPU 之間的數據同步的 SyncedMemory 類,而這也是後面網絡參數,梯度數據傳輸的基礎。
主要包含了:

  • CaffeMallocHost 內存申請函數
  • CaffeFreeHost 內存釋放函數
  • SyncedMemory 類

1. CaffeMallocHost 函數

inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
#ifndef CPU_ONLY // 如果是 GPU 模式就調用 cudaMallocHost 
  if (Caffe::mode() == Caffe::GPU) {
	CUDA_CHECK(cudaMallocHost(ptr, size));
	*use_cuda = true;
	return;
  }
#endif
#ifdef USE_MKL // 使用 mkl 庫則使用 mkl_malloc
  *ptr = mkl_malloc(size ? size:1, 64);
#else
  *ptr = malloc(size); // 正常 CPU 就直接使用 malloc
#endif
  *use_cuda = false;
  CHECK(*ptr) << "host allocation of size " << size << " failed";
}

可以看到上述函數根據當前環境的不同使用了不同的內存申請函數來完成爲數據申請內存,CPU 模式下就直接用了 malloc 函數。

2. CaffeFreeHost 函數

與第 1 部分對應,我們申請了內存就要負責將資源釋放掉。

inline void CaffeFreeHost(void* ptr, bool use_cuda) {
#ifndef CPU_ONLY // GPU 模式
  if (use_cuda) {
	CUDA_CHECK(cudaFreeHost(ptr));
	return;
  }
#endif
#ifdef USE_MKL  // 使用 mkl 庫
  mkl_free(ptr);
#else
  free(ptr);  // CPU 模式
#endif
}

3. SyncedMemory 類

主要成員變量 主要成員函數
void* cpu_ptr_; void to_cpu();
void* gpu_ptr_; void to_gpu();
bool own_cpu_data_; void set_cpu_data(void* data);
bool own_gpu_data_; void set_gpu_data(void* data);
SyncedHead head_; ——

這個類主要的作用就是同步 CPU 與 GPU 之間的數據,主要是通過上面幾個主要的成員變量來進行控制,SyncedHead head_ 是一個枚舉類型。

  enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };

狀態機變量,表示4種狀態:未初始化、CPU數據有效、GPU數據有效、已同步。 以 to_cpu() 成員函數爲例。

inline void SyncedMemory::to_cpu() {
  check_device();
  switch (head_) {
  case UNINITIALIZED:  // 如果未分配過內存(構造函數後就是這個狀態)
    CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);  // to_CPU時爲CPU分配內存
    caffe_memset(size_, 0, cpu_ptr_); // 數據清零
    head_ = HEAD_AT_CPU;  // 指示CPU更新了數據
    own_cpu_data_ = true;
    break;
  case HEAD_AT_GPU: // 如果GPU側更新過數據,則同步到CPU
#ifndef CPU_ONLY
    if (cpu_ptr_ == NULL) { // 如果CPU側沒分配過內存,分配內存
      CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
      own_cpu_data_ = true;
    }
    caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_);  // 數據同步
    head_ = SYNCED;  // 指示CPU和GPU數據已同步一致
#else
    NO_GPU;
#endif
    break;
  case HEAD_AT_CPU:  // 如果CPU數據是最新的,不操作
  case SYNCED:  // 如果CPU和GPU數據都是最新的,不操作
    break;
  }
}

可以看到如果當前數據的狀態是未初始化,則通過調用內存分配函數對數據進行內存申請,並置成 0,最後賦到 cpu_ptr_ 返回,這裏如果是經由 to_cpu()進行初始化的數據,有一個標誌位 own_cpu_data_ 會被設成 true,表示當前擁有數據的所有權,這個所有權表示我們申請它同樣要負責釋放它。以當前數據在 CPU 上爲例,如果當前數據需要同步到 GPU 上,則需要使用 caffe_gpu_memcpy() 來完成數據同步工作並表示數據同步已完成。

void SyncedMemory::set_cpu_data(void* data) {
  check_device();
  CHECK(data);
  if (own_cpu_data_) {  // 如果自己分配過內存,先釋放,換外部指定數據
    CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
  }
  cpu_ptr_ = data;  // 直接指向外部數據
  head_ = HEAD_AT_CPU;  // 指示CPU側更新了數據
  own_cpu_data_ = false;   // 指示數據來源於外部
}

主要體現在 set_cpu_data() 這個成員函數中,當我們從外部輸入data這裏的數據可以是網絡參數也可以是參數梯度,如果當前我們擁有數據的所有權就需要將當前申請的數據釋放掉,然後再將 cpu_ptr_ 指向輸入的數據 data,完成數據設置,這時標誌位 own_cpu_data_ 會被設成 false

其他的成員函數還有 mutable_cpu_data() 這個也是直接調用了 to_cpu() 函數,在後面的 blob.cpp 中會有調用。

blob.hpp/cpp

談到 Blob,它是 caffe 框架中數據的傳輸格式,是一個四維矩陣,分別對應:N,C,H,W,即 批大小,通道數, 特徵圖的高, 特徵圖的寬,並且提供了數據 data 和對應梯度 diff 的操作方式,爲後面網絡的訓練提供基礎。

本文件主要是定義了 Blob 的類,其與 SyncedMemory 類構成一種委託的關係,即 blob 類裏包含了 SyncedMemory 類的指針。

  shared_ptr<SyncedMemory> data_;  // 存放指向 data 的指針
  shared_ptr<SyncedMemory> diff_; // 存放指向 diff 的指針
  shared_ptr<SyncedMemory> shape_data_; // 存放維度信息的指針
  vector<int> shape_; // 形狀信息
  int count_; // 存放有效元素數目
  int capacity_; // 存放 Blob 容器的容量信息

在 blob 類構造函數中主要就是在內部先初始化容量信息 capacity_,然後調用了 void Reshape(); 函數,其功能是改變 blob 的維度,如果有需要需要重新申請更大的內存。根據輸入格式的不同, void Reshape() 函數重載了3種,分別是:

void Reshape(const int num, const int channels, const int height,const int width);  
void Reshape(const BlobShape& shape);
void ReshapeLike(const Blob& other);

不管哪種形式最後都是統一轉化並調用:

void Reshape(const vector<int>& shape); 

來完成上述功能,下面是 Reshape 的函數定義。

template <typename Dtype>
void Blob<Dtype>::Reshape(const vector<int>& shape) {   
  CHECK_LE(shape.size(), kMaxBlobAxes); // 先檢查輸入數據維數有沒有大於最大維數,合法是4
  count_ = 1; // 用於計算元素總數 count_=num*channels*height*width, 初始輸入是 1x3x48x48
  shape_.resize(shape.size()); // 成員變量維度也被重置,這裏是 vector 的 resize,shape_初始是 0 ,運行後變成 [0,0,0,0]
  if (!shape_data_ || shape_data_->size() < shape.size() * sizeof(int)) {  // sizeof(int) 爲4, shape.size() 爲4
    shape_data_.reset(new SyncedMemory(shape.size() * sizeof(int)));  // shape_data_ 的 size_ 爲 16
  }
  int* shape_data = static_cast<int*>(shape_data_->mutable_cpu_data()); // 指針格式轉化
  for (int i = 0; i < shape.size(); ++i) {
    CHECK_GE(shape[i], 0); // 保證每一個維度尺寸都 >=0
    if (count_ != 0) {
		// 保證 count_ 不溢出,應爲下一次 count_ 的數量就是 shape[i] * count_
      CHECK_LE(shape[i], INT_MAX / count_) << "blob size exceeds INT_MAX";
    }
    count_ *= shape[i]; // count_ 累乘
    shape_[i] = shape[i]; // 爲成員變量賦值 shape_[i] 分別爲 1  3  48  48
    shape_data[i] = shape[i];
  }
  if (count_ > capacity_) { //如果當前成員大於分配的空間容量,初始的 capacity_ 爲 0
    capacity_ = count_; // 擴容,重新分配 data_ 和 diff_ 的空間
    data_.reset(new SyncedMemory(capacity_ * sizeof(Dtype)));  // data_ 要分配 6912 * sizeof(Dtype) 個字節的內存大小
    diff_.reset(new SyncedMemory(capacity_ * sizeof(Dtype)));  // diff_ 要分配 6912 * sizeof(Dtype) 個字節的內存大小
  }
}

可以看到是根據輸入數據得 shape 信息與當前分配的空間容量做比較,然後重新分配對應的空間來存儲數據。基於上面的構造函數,然後在 Blob 類中,使用 Update() 來更新網絡參數,即 data=data-diff 來更新。

// Update() 函數用於網絡參數 Blob 的更新
template <typename Dtype>
void Blob<Dtype>::Update() {
  // We will perform update based on where the data is located.
  switch (data_->head()) {  // data 在哪就在哪更新
  case SyncedMemory::HEAD_AT_CPU:  // data 位於 CPU 端
    // perform computation on CPU
	  // 執行在 CPU 上的計算 data_[i] = data_[i]-diff[i]
    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: // data 位於 GPU 端,或者 CPU/GPU已同步
  case SyncedMemory::SYNCED:
#ifndef CPU_ONLY
    // perform computation on GPU
    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.";
  }
}

這裏的 caffe_axpy() 則是調用了標準 cblas 加速庫的東西,同理,還有其他計算 data L1,L2範數的函數定義,以及它們乘以一個標量的函數,均是通過 data_->head() 的情況來調用加速庫進行相關計算。
餘下的就是一些對 blob 數據結構的查詢功能,例如某一維的尺寸,對輸入 index 的轉化來查詢到想對應的數據,還有就是數據的 Offset, 根據頭指針及維度信息來查詢到維度對應的實際存儲信息的位置等。

其他的數據操作,訪問 cpu 數據啊,讀寫 cpu 數據啊,都是直接調用的 SyncedMemory 類。就這樣構建起了整個 caffe 框架的數據存儲結構!

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章