【Caffe代碼解析】SyncedMemory

【Caffe代碼解析】SyncedMemory

功能：

Caffe的底層數據的切換（cpu模式和gpu模式），需要用到內存同步模塊。

源碼：頭文件

#ifndef CAFFE_SYNCEDMEM_HPP_
#define CAFFE_SYNCEDMEM_HPP_

#include <cstdlib>

#include "caffe/common.hpp" #include "caffe/util/math_functions.hpp" namespace caffe { inline void CaffeMallocHost(void** ptr, size_t size) { *ptr = malloc(size); CHECK(*ptr) << "host allocation of size " << size << " failed"; } inline void CaffeFreeHost(void* ptr) { free(ptr); } /** * @brief Manages memory allocation and synchronization between the host (CPU) * and device (GPU). * * TODO(dox): more thorough description. */ class SyncedMemory { public: SyncedMemory() : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED), own_cpu_data_(false) {} explicit SyncedMemory(size_t size) : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED), own_cpu_data_(false) {} ~SyncedMemory(); const void* cpu_data();//獲取cpu數據，返回void * 指針 void set_cpu_data(void* data);//用一個void * 指針修改指針 const void* gpu_data();//獲取gpu數據，返回void * 指針 void* mutable_cpu_data();//獲取可以更改cpu數據，返回void * 指針 void* mutable_gpu_data();//獲取可以更改gpu數據，返回void * 指針 enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };//enum枚舉值 SyncedHead head() { return head_; }//獲得枚舉值 size_t size() { return size_; }//獲得數據大小 private: void to_cpu();//轉為cpu模式 void to_gpu(); //轉為gpu模式 void* cpu_ptr_;//指向cpu的指針 void* gpu_ptr_;//指向gpu的指指針 size_t size_; //大小 SyncedHead head_; //數據存放的位置，枚舉值之一 bool own_cpu_data_;//是否有cpu數據 DISABLE_COPY_AND_ASSIGN(SyncedMemory); }; // class SyncedMemory } // namespace caffe #endif // CAFFE_SYNCEDMEM_HPP_

實現文件：

#include <cstring> #include "caffe/common.hpp" #include "caffe/syncedmem.hpp" #include "caffe/util/math_functions.hpp" namespace caffe { //析構函數，調用caffe函數來釋放空間 SyncedMemory::~SyncedMemory() { if (cpu_ptr_ && own_cpu_data_) { CaffeFreeHost(cpu_ptr_); } //如果有gpu數據，也進行釋放。 #ifndef CPU_ONLY if (gpu_ptr_) { CUDA_CHECK(cudaFree(gpu_ptr_)); } #endif // CPU_ONLY } // 根據head信息，選擇：1， 分類cpu空間，2.拷貝gpu數據值cpu inline void SyncedMemory::to_cpu() { switch (head_) { case UNINITIALIZED: CaffeMallocHost(&cpu_ptr_, size_); caffe_memset(size_, 0, cpu_ptr_); head_ = HEAD_AT_CPU; own_cpu_data_ = true; break; case HEAD_AT_GPU: #ifndef CPU_ONLY if (cpu_ptr_ == NULL) { CaffeMallocHost(&cpu_ptr_, size_); own_cpu_data_ = true; } caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_); head_ = SYNCED; #else NO_GPU; #endif break; case HEAD_AT_CPU: case SYNCED: break; } } // 根據head信息，選擇：1， 分類gpu空間，2.拷貝cpu數據值gpu inline void SyncedMemory::to_gpu() { #ifndef CPU_ONLY switch (head_) { case UNINITIALIZED: CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_)); caffe_gpu_memset(size_, 0, gpu_ptr_); head_ = HEAD_AT_GPU; break; case HEAD_AT_CPU: if (gpu_ptr_ == NULL) { CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_)); } caffe_gpu_memcpy(size_, cpu_ptr_, gpu_ptr_); head_ = SYNCED; break; case HEAD_AT_GPU: case SYNCED: break; } #else NO_GPU; #endif } //返回cpu指針 void * 類型 const void* SyncedMemory::cpu_data() { to_cpu(); return (const void*)cpu_ptr_; } // 設置cpu數據，利用另外一個指針的數據來初始化 void SyncedMemory::set_cpu_data(void* data) { CHECK(data); if (own_cpu_data_) { CaffeFreeHost(cpu_ptr_); } cpu_ptr_ = data;//直接重置指針， head_ = HEAD_AT_CPU; own_cpu_data_ = false;//設false } //獲得gpu指針 const void* SyncedMemory::gpu_data() { #ifndef CPU_ONLY to_gpu(); return (const void*)gpu_ptr_; #else NO_GPU; #endif } //獲得可更改的cpu指針 void* SyncedMemory::mutable_cpu_data() { to_cpu(); head_ = HEAD_AT_CPU; return cpu_ptr_; } //獲得可更改的gpu指針 void* SyncedMemory::mutable_gpu_data() { #ifndef CPU_ONLY to_gpu(); head_ = HEAD_AT_GPU; return gpu_ptr_; #else NO_GPU; #endif } } // namespace caffe