Caffe源碼解讀：syncedmem類

內(nèi)存同步(syncedmem)類的作用在于管理主機(jī)（CPU）和設(shè)備（GPU）之間的內(nèi)存分配和數(shù)據(jù)同步，封裝了二者之間的交互操作。

這個(gè)類沒有對(duì)應(yīng)的ProtoBuffer描述，所以直接看./include/caffe/syncedmem.cpp文件：

#ifndef CAFFE_SYNCEDMEM_HPP_
#define CAFFE_SYNCEDMEM_HPP_

#include <cstdlib>

#include "caffe/common.hpp"

namespace caffe {

// 以頁鎖定方式分配內(nèi)存，在單GPU時(shí)提示不明顯，多GPU提升很多
// malloc分配標(biāo)準(zhǔn)可分頁的主機(jī)內(nèi)存，操作系統(tǒng)可能會(huì)將這種內(nèi)存分頁或者交換到磁盤上，需要的時(shí)候調(diào)回內(nèi)存，這樣可能會(huì)增加運(yùn)行時(shí)間
// cudaMallocHost分配頁鎖定的主機(jī)內(nèi)存，操作系統(tǒng)不會(huì)對(duì)這塊內(nèi)存分頁或者交換到磁盤上，可以節(jié)省時(shí)間
inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
#ifndef CPU_ONLY
  if (Caffe::mode() == Caffe::GPU) {
    CUDA_CHECK(cudaMallocHost(ptr, size)); // 分配顯存并校驗(yàn)是否分配成功
    *use_cuda = true;
    return;
  }
#endif
  *ptr = malloc(size); // CPU模式下分配內(nèi)存
  *use_cuda = false;
  CHECK(*ptr) << "host allocation of size " << size << " failed";
}

// 和上面函數(shù)相對(duì)應(yīng)的內(nèi)存釋放
inline void CaffeFreeHost(void* ptr, bool use_cuda) {
#ifndef CPU_ONLY
  if (use_cuda) {
    //如果分配內(nèi)存用的是cudaMallocHost分配，即use_cuda為真，cpu中的數(shù)據(jù)也可以用cudaMallocHost分配內(nèi)存  
    CUDA_CHECK(cudaFreeHost(ptr));  
    return;
  }
#endif
  free(ptr);  // 用的malloc分配的內(nèi)存
}

// 負(fù)責(zé)內(nèi)存分配和設(shè)備同步
class SyncedMemory {
 public:
  SyncedMemory()     // 默認(rèn)構(gòu)造函數(shù) 
      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
        own_cpu_data_(false), cpu_malloc_use_cuda_(false), own_gpu_data_(false),
        gpu_device_(-1) {}
  explicit SyncedMemory(size_t size)  // 顯式構(gòu)造函數(shù)
      : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED),
        own_cpu_data_(false), cpu_malloc_use_cuda_(false), own_gpu_data_(false),
        gpu_device_(-1) {}
  ~SyncedMemory();

  // 對(duì)CPU，GPU數(shù)據(jù)的讀寫，不贅述。為什么沒有set_cpu_diff() ???
  const void* cpu_data();
  void set_cpu_data(void* data);
  const void* gpu_data();
  void set_gpu_data(void* data);
  void* mutable_cpu_data();
  void* mutable_gpu_data();

  // 共享內(nèi)存的4種狀態(tài)：未初始化，CPU數(shù)據(jù)有效，GPU數(shù)據(jù)有效，已同步
  enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
  // 返回當(dāng)前共享內(nèi)存的狀態(tài)
  SyncedHead head() { return head_; }
  // 返回存儲(chǔ)空間的尺寸 = 元素?cái)?shù) * 單個(gè)元素所占字節(jié)數(shù)
  size_t size() { return size_; }

#ifndef CPU_ONLY
  void async_gpu_push(const cudaStream_t& stream);
#endif

 private:
  void to_cpu();  // 數(shù)據(jù)同步至cpu
  void to_gpu();  // 數(shù)據(jù)同步至gpu
  void* cpu_ptr_;  // cpu中數(shù)據(jù)的指針
  void* gpu_ptr_;   // gpu中數(shù)據(jù)的指針
  size_t size_;   // 存儲(chǔ)空間的大小
  SyncedHead head_;  // 共享內(nèi)存的狀態(tài)
  bool own_cpu_data_;  // cpu擁有數(shù)據(jù)所有權(quán)
  bool cpu_malloc_use_cuda_;   // 分配cpu內(nèi)存是否用cudaMallocHost()分配。
  bool own_gpu_data_;  // gpu擁有數(shù)據(jù)所有權(quán)
  int gpu_device_;  // gpu設(shè)備號(hào)

  // 禁用拷貝構(gòu)造以及賦值運(yùn)算符
  // 使用grep可以查到，該宏定義在common.hpp第35行
  // 該宏就是把拷貝構(gòu)造和賦值運(yùn)算符設(shè)置為private而已
  DISABLE_COPY_AND_ASSIGN(SyncedMemory); 
};  // class SyncedMemory

}  // namespace caffe
#endif  // CAFFE_SYNCEDMEM_HPP_

這個(gè)類比Blob簡單多了，下面看對(duì)應(yīng)的./src/caffe/syncedmem.cpp文件：

#include "caffe/common.hpp"
#include "caffe/syncedmem.hpp"
#include "caffe/util/math_functions.hpp"

namespace caffe {

SyncedMemory::~SyncedMemory() {
 // 如果cpu擁有數(shù)據(jù)所有權(quán)
 if (cpu_ptr_ && own_cpu_data_) {
   CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
 }
 // 如果數(shù)據(jù)在gpu上
#ifndef CPU_ONLY
 if (gpu_ptr_ && own_gpu_data_) {
   int initial_device;
   cudaGetDevice(&initial_device);  // 獲取使用的gpu設(shè)備號(hào)
   if (gpu_device_ != -1) {
     CUDA_CHECK(cudaSetDevice(gpu_device_));
   }
   CUDA_CHECK(cudaFree(gpu_ptr_));
   cudaSetDevice(initial_device);
 }
#endif  // CPU_ONLY
}

// 數(shù)據(jù)同步到cpu上
inline void SyncedMemory::to_cpu() {
 switch (head_) {
 case UNINITIALIZED:  // 如果是未初始化狀態(tài)，只需分配下內(nèi)存就行了
   CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
   caffe_memset(size_, 0, cpu_ptr_);  // 定義在math_functions.hpp第42行，調(diào)用了memset
   head_ = HEAD_AT_CPU;  // 內(nèi)存狀態(tài)改為cpu擁有所有權(quán)
   own_cpu_data_ = true;
   break;
 case HEAD_AT_GPU:  // GPU擁有數(shù)據(jù)所有權(quán)
#ifndef CPU_ONLY
   if (cpu_ptr_ == NULL) {
     CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);// 要內(nèi)存
     own_cpu_data_ = true;
   }
   caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_); // 數(shù)據(jù)復(fù)制
   head_ = SYNCED;
#else
   NO_GPU;
#endif
   break;
 // 數(shù)據(jù)已經(jīng)為cpu擁有所有權(quán)或者在內(nèi)存共享狀態(tài)，則什么都不管
 case HEAD_AT_CPU:
 case SYNCED:
   break;
 }
}

// 原理同上
inline void SyncedMemory::to_gpu() {
#ifndef CPU_ONLY
 switch (head_) {
 case UNINITIALIZED:
   CUDA_CHECK(cudaGetDevice(&gpu_device_));
   CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
   caffe_gpu_memset(size_, 0, gpu_ptr_);
   head_ = HEAD_AT_GPU;
   own_gpu_data_ = true;
   break;
 case HEAD_AT_CPU:
   if (gpu_ptr_ == NULL) {
     CUDA_CHECK(cudaGetDevice(&gpu_device_));
     CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_)); // 要一塊內(nèi)存
     own_gpu_data_ = true;
   }
   caffe_gpu_memcpy(size_, cpu_ptr_, gpu_ptr_); //數(shù)據(jù)拷貝，調(diào)用了cudaMemcpy函數(shù) 
   head_ = SYNCED;
   break;
 case HEAD_AT_GPU:
 case SYNCED:
   break;
 }
#else
 NO_GPU;
#endif
}

// 獲取cpu中的數(shù)據(jù)，只讀
const void* SyncedMemory::cpu_data() {
 to_cpu();
 return (const void*)cpu_ptr_;
}
// 設(shè)置獲取cpu中的數(shù)據(jù)
void SyncedMemory::set_cpu_data(void* data) {
 CHECK(data);
 if (own_cpu_data_) {
   // 調(diào)用這個(gè)函數(shù)的時(shí)候，如果cpu內(nèi)有數(shù)據(jù)會(huì)被直接清空，要注意
   CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_); 
 }
 cpu_ptr_ = data;
 head_ = HEAD_AT_CPU;
 own_cpu_data_ = false;
}
// 獲取gpu中的數(shù)據(jù)，只讀
const void* SyncedMemory::gpu_data() {
#ifndef CPU_ONLY
 to_gpu();
 return (const void*)gpu_ptr_;
#else
 NO_GPU;
 return NULL;
#endif
}
// 設(shè)置gpu中的數(shù)據(jù)
void SyncedMemory::set_gpu_data(void* data) {
#ifndef CPU_ONLY
 CHECK(data);
 if (own_gpu_data_) {
   int initial_device;
   cudaGetDevice(&initial_device);
   if (gpu_device_ != -1) {
     CUDA_CHECK(cudaSetDevice(gpu_device_));
   }
 // 調(diào)用這個(gè)函數(shù)的時(shí)候，如果gpu內(nèi)有數(shù)據(jù)會(huì)被直接清空，要注意
   CUDA_CHECK(cudaFree(gpu_ptr_));
   cudaSetDevice(initial_device);
 }
 gpu_ptr_ = data;
 head_ = HEAD_AT_GPU;
 own_gpu_data_ = false;
#else
 NO_GPU;
#endif
}
// 讀寫獲取cpu數(shù)據(jù)
void* SyncedMemory::mutable_cpu_data() {
 to_cpu();
 head_ = HEAD_AT_CPU;
 return cpu_ptr_;
}
// 讀寫獲取gpu數(shù)據(jù)
void* SyncedMemory::mutable_gpu_data() {
#ifndef CPU_ONLY
 to_gpu();
 head_ = HEAD_AT_GPU;
 return gpu_ptr_;
#else
 NO_GPU;
 return NULL;
#endif
}

// cuda中的流同步，這里傳入一個(gè)異步流，在計(jì)算的時(shí)候向GPU復(fù)制數(shù)據(jù)。
#ifndef CPU_ONLY
void SyncedMemory::async_gpu_push(const cudaStream_t& stream) {
 CHECK(head_ == HEAD_AT_CPU);
 if (gpu_ptr_ == NULL) {
   CUDA_CHECK(cudaGetDevice(&gpu_device_));
   CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
   own_gpu_data_ = true;
 }
 const cudaMemcpyKind put = cudaMemcpyHostToDevice;
 CUDA_CHECK(cudaMemcpyAsync(gpu_ptr_, cpu_ptr_, size_, put, stream));
 // Assume caller will synchronize on the stream before use
 head_ = SYNCED;
}
#endif

}  // namespace caffe

syncedmem類比較簡單，主要是完成cpu和gpu之間的數(shù)據(jù)交互問題~

參考資料

  1. 《21天實(shí)戰(zhàn)caffe》
  2. （介紹了流同步）http://blog.csdn.net/sinat_22336563/article/details/68496919
最后編輯于 ：
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成，瀏覽時(shí)請結(jié)合常識(shí)與多方信息審慎甄別。
平臺(tái)聲明：文章內(nèi)容（如有圖片或視頻亦包括在內(nèi)）由作者上傳并發(fā)布，文章內(nèi)容僅代表作者本人觀點(diǎn)，簡書系信息發(fā)布平臺(tái)，僅提供信息存儲(chǔ)服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容