caffe源码分析-SyncedMemory

bear_fish

发布于 2019-02-25 11:49:36

5680

发布于 2019-02-25 11:49:36

文章被收录于专栏：用户2442861的专栏

本文主要分析**caffe**中**Blob**内存管理类**SyncedMemory**，主要内容包括：

SyncedMemory和Blob的关系
SyncedMemory中的方法，如内存的分配、释放
SyncedMemory中内存的申请，是在数据访问时才分配而不是立马分配(通过enum SyncedHead状态实现)

`SyncedMemory`和`Blob`的关系

Blob中的主要数据成员如下，实际是在SyncedMemory上做了一层包装：

template<typename Dtype>
class Blob{    
   protected:
        shared_ptr<SyncedMemory> data_; //存储前向传递数据
        shared_ptr<SyncedMemory> diff_; //存储反向传递梯度
        vector<int> shape_;  //参数维度
        int count_; //Blob存储的元素个数（shape_所有元素乘积）
        int capacity_;//当前Blob的元素个数（控制动态分配）
};

`SyncedMemory`

SyncedMemory 的主要数据成员如下：

class SyncedMemory {
   private:
        void* cpu_ptr_;//内存指针
        void* gpu_ptr_;//显存指针
        size_t size_;  //数据大小
        SyncedHead head_;//当前数据状态
        bool own_cpu_data_;
        bool cpu_malloc_use_cuda_;
        bool own_gpu_data_;
        int gpu_device_;
};

SyncedMemory屏蔽了代码对不同硬件设备的内存分配的感知，同时隐藏了CPU和GPU之间的同步过程。
SyncedMemory采用“lazy”的模式，就是内存的实际申请时机是在第一次使用时进行的(通过枚举状态)。

首先看如下两个函数cpu_data， gpu_data获取cpu，gpu数据指针:

const void* SyncedMemory::cpu_data() {
    to_cpu();   // 首先完成数据同步，第一次访问时会申请存储空间
    return (const void*)cpu_ptr_;
}
const void* SyncedMemory::gpu_data() {
    #ifndef CPU_ONLY
            to_gpu();
            return (const void*)gpu_ptr_;
    #else
            NO_GPU;
      return NULL;
    #endif
    return nullptr;
}

以cpu_data为例, 如果数据尚未分配(通过如下枚举判断)，则分配数据，如果已经分配则什么也不做：

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

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
//     ........
            ;
#else
            NO_GPU;
#endif
            break;
        case HEAD_AT_CPU:
        case SYNCED:
            break;
    }
}

SyncedMemory有如下接口，获取cpu中的数据指针(gpu同理)：

//获取CPU数据指针，不能改变数据内容
const void* cpu_data();
//获取CPU数据指针，可以改变数据内容
void* mutable_cpu_data();

Blob中调用:

// 调用SyncedMemory的数据访问函数cpu_data(),并返回内存指针
template <typename Dtype>
const Dtype* Blob<Dtype>::cpu_data() const {
    CHECK(data_);
    return (const Dtype*)data_->cpu_data();
}
template <typename Dtype>
Dtype* Blob<Dtype>::mutable_cpu_data() {
    CHECK(data_);
    return static_cast<Dtype*>(data_->mutable_cpu_data());
}

内存分配与释放

内存分配与释放由两个(不属于SyncedMemory类)的内联函数完成. 代码简单直观: 如果是CPU模式, 那么调用malloc和free来申请/释放内存, 否则调用CUDA的cudaMallocHost和cudaFreeHost来申请/释放显存.

inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
    *ptr = malloc(size);
    *use_cuda = false;
    CHECK(*ptr) << "host allocation of size " << size << " failed";
}

inline void CaffeFreeHost(void* ptr, bool use_cuda) {
    free(ptr);
}