SyncedMemory類定義在syncedmem.hpp/cpp里, 負(fù)責(zé)caffe底層的內(nèi)存管理.
內(nèi)存分配與釋放
內(nèi)存分配與釋放由兩個(gè)(不屬于SyncedMemory類)的內(nèi)聯(lián)函數(shù)完成. 代碼簡單直觀: 如果是CPU模式, 那么調(diào)用malloc和free來申請(qǐng)/釋放內(nèi)存, 否則調(diào)用CUDA的cudaMallocHost和cudaFreeHost來申請(qǐng)/釋放顯存.
// ------ 分配內(nèi)存 ------
inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
#ifndef CPU_ONLY
if (Caffe::mode() == Caffe::GPU) {
CUDA_CHECK(cudaMallocHost(ptr, size));
*use_cuda = true;
return;
}
#endif
*ptr = malloc(size);
*use_cuda = false;
CHECK(*ptr) << "host allocation of size " << size << " failed";
}
// ------ 釋放內(nèi)存 ------
inline void CaffeFreeHost(void* ptr, bool use_cuda) {
#ifndef CPU_ONLY
if (use_cuda) {
CUDA_CHECK(cudaFreeHost(ptr));
return;
}
#endif
free(ptr);
}
類成員變量
void* cpu_ptr_; // cpu 內(nèi)存地址
void* gpu_ptr_; // gpu 內(nèi)存地址
size_t size_; // 數(shù)據(jù)大小
SyncedHead head_; // 當(dāng)前數(shù)據(jù)同步狀態(tài)
bool own_cpu_data_; // 是否是自己的cpu data? (例如set_cpu_data就是false)
bool cpu_malloc_use_cuda_;
bool own_gpu_data_; // 是否已經(jīng)申請(qǐng)gpu內(nèi)存空間
int gpu_device_; //
值得稍加注意的是SyncedHead head_. 該變量的作用會(huì)在數(shù)據(jù)同步部分說明.
get and set 方法
cpu_data, gpu_data或者mutable_cpu_data, mutable_gpu_data方法返回cpu或者gpu內(nèi)存指針, 前者是const void*, 不可對(duì)返回內(nèi)存進(jìn)行修改; 后者為void*, 可以修改.
set方法比較特別, 方法參數(shù)是指向另一段內(nèi)存空間的地址:
void SyncedMemory::set_cpu_data(void* data) {
CHECK(data);
if (own_cpu_data_) {
CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
}
cpu_ptr_ = data;
head_ = HEAD_AT_CPU;
own_cpu_data_ = false;
}
該函數(shù)首先釋放自己申請(qǐng)的內(nèi)存空間, 然后直接指向參數(shù)傳入的內(nèi)存空間 (并不是重新申請(qǐng)空間, 并copy數(shù)據(jù)). 最后將 own_cpu_data_設(shè)置為false, 表示外來數(shù)據(jù)(?).
保持?jǐn)?shù)據(jù)同步
在調(diào)用cpu_data或者gpu_data方法時(shí), 需要確保cpu, gpu數(shù)據(jù)內(nèi)容是一致的. 這里用到了前面提到的枚舉類型來記錄當(dāng)前同步狀態(tài)
enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
以to_cpu()方法為例: 檢查head_所處狀態(tài), 若UNINITIALIZED, 則分配內(nèi)存空間(置0); 若HEAD_AT_GPU, 則需要從GPU內(nèi)存同步數(shù)據(jù)到CPU; HEAD_AT_CPU, 則說明目前最新的數(shù)據(jù)是在CPU的, 無須進(jìn)行任何操作 (雖然并不知道GPU的數(shù)據(jù)是否和CPU一致, 因?yàn)楫?dāng)前我們并不關(guān)心GPU數(shù)據(jù)); 若SYNCED, 則CPU/GPU數(shù)據(jù)一致, 無須進(jìn)行任何操作.
inline void SyncedMemory::to_cpu() {
switch (head_) {
case UNINITIALIZED:
CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
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_, &cpu_malloc_use_cuda_);
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;
}
}
