Caffe源码解读:syncedmem类

作者: Mordekaiser | 来源:发表于2017-07-24 17:13 被阅读103次

    内存同步(syncedmem)类的作用在于管理主机(CPU)和设备(GPU)之间的内存分配和数据同步,封装了二者之间的交互操作。

    这个类没有对应的ProtoBuffer描述,所以直接看./include/caffe/syncedmem.cpp文件:

    #ifndef CAFFE_SYNCEDMEM_HPP_
    #define CAFFE_SYNCEDMEM_HPP_
    
    #include <cstdlib>
    
    #include "caffe/common.hpp"
    
    namespace caffe {
    
    // 以页锁定方式分配内存,在单GPU时提示不明显,多GPU提升很多
    // malloc分配标准可分页的主机内存,操作系统可能会将这种内存分页或者交换到磁盘上,需要的时候调回内存,这样可能会增加运行时间
    // cudaMallocHost分配页锁定的主机内存,操作系统不会对这块内存分页或者交换到磁盘上,可以节省时间
    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); // CPU模式下分配内存
      *use_cuda = false;
      CHECK(*ptr) << "host allocation of size " << size << " failed";
    }
    
    // 和上面函数相对应的内存释放
    inline void CaffeFreeHost(void* ptr, bool use_cuda) {
    #ifndef CPU_ONLY
      if (use_cuda) {
        //如果分配内存用的是cudaMallocHost分配,即use_cuda为真,cpu中的数据也可以用cudaMallocHost分配内存  
        CUDA_CHECK(cudaFreeHost(ptr));  
        return;
      }
    #endif
      free(ptr);  // 用的malloc分配的内存
    }
    
    // 负责内存分配和设备同步
    class SyncedMemory {
     public:
      SyncedMemory()     // 默认构造函数 
          : 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)  // 显式构造函数
          : 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();
    
      // 对CPU,GPU数据的读写,不赘述。为什么没有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();
    
      // 共享内存的4种状态:未初始化,CPU数据有效,GPU数据有效,已同步
      enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
      // 返回当前共享内存的状态
      SyncedHead head() { return head_; }
      // 返回存储空间的尺寸 = 元素数 * 单个元素所占字节数
      size_t size() { return size_; }
    
    #ifndef CPU_ONLY
      void async_gpu_push(const cudaStream_t& stream);
    #endif
    
     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拥有数据所有权
      bool cpu_malloc_use_cuda_;   // 分配cpu内存是否用cudaMallocHost()分配。
      bool own_gpu_data_;  // gpu拥有数据所有权
      int gpu_device_;  // gpu设备号
    
      // 禁用拷贝构造以及赋值运算符
      // 使用grep可以查到,该宏定义在common.hpp第35行
      // 该宏就是把拷贝构造和赋值运算符设置为private而已
      DISABLE_COPY_AND_ASSIGN(SyncedMemory); 
    };  // class SyncedMemory
    
    }  // namespace caffe
    #endif  // CAFFE_SYNCEDMEM_HPP_
    

    这个类比Blob简单多了,下面看对应的./src/caffe/syncedmem.cpp文件:

    #include "caffe/common.hpp"
    #include "caffe/syncedmem.hpp"
    #include "caffe/util/math_functions.hpp"
    
    namespace caffe {
    
    SyncedMemory::~SyncedMemory() {
     // 如果cpu拥有数据所有权
     if (cpu_ptr_ && own_cpu_data_) {
       CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
     }
     // 如果数据在gpu上
    #ifndef CPU_ONLY
     if (gpu_ptr_ && own_gpu_data_) {
       int initial_device;
       cudaGetDevice(&initial_device);  // 获取使用的gpu设备号
       if (gpu_device_ != -1) {
         CUDA_CHECK(cudaSetDevice(gpu_device_));
       }
       CUDA_CHECK(cudaFree(gpu_ptr_));
       cudaSetDevice(initial_device);
     }
    #endif  // CPU_ONLY
    }
    
    // 数据同步到cpu上
    inline void SyncedMemory::to_cpu() {
     switch (head_) {
     case UNINITIALIZED:  // 如果是未初始化状态,只需分配下内存就行了
       CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
       caffe_memset(size_, 0, cpu_ptr_);  // 定义在math_functions.hpp第42行,调用了memset
       head_ = HEAD_AT_CPU;  // 内存状态改为cpu拥有所有权
       own_cpu_data_ = true;
       break;
     case HEAD_AT_GPU:  // 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;
     // 数据已经为cpu拥有所有权或者在内存共享状态,则什么都不管
     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_)); // 要一块内存
         own_gpu_data_ = true;
       }
       caffe_gpu_memcpy(size_, cpu_ptr_, gpu_ptr_); //数据拷贝,调用了cudaMemcpy函数 
       head_ = SYNCED;
       break;
     case HEAD_AT_GPU:
     case SYNCED:
       break;
     }
    #else
     NO_GPU;
    #endif
    }
    
    // 获取cpu中的数据,只读
    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_) {
       // 调用这个函数的时候,如果cpu内有数据会被直接清空,要注意
       CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_); 
     }
     cpu_ptr_ = data;
     head_ = HEAD_AT_CPU;
     own_cpu_data_ = false;
    }
    // 获取gpu中的数据,只读
    const void* SyncedMemory::gpu_data() {
    #ifndef CPU_ONLY
     to_gpu();
     return (const void*)gpu_ptr_;
    #else
     NO_GPU;
     return NULL;
    #endif
    }
    // 设置gpu中的数据
    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_));
       }
     // 调用这个函数的时候,如果gpu内有数据会被直接清空,要注意
       CUDA_CHECK(cudaFree(gpu_ptr_));
       cudaSetDevice(initial_device);
     }
     gpu_ptr_ = data;
     head_ = HEAD_AT_GPU;
     own_gpu_data_ = false;
    #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;
     return NULL;
    #endif
    }
    
    // cuda中的流同步,这里传入一个异步流,在计算的时候向GPU复制数据。
    #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之间的数据交互问题~

    参考资料

    1. 《21天实战caffe》
    2. (介绍了流同步)http://blog.csdn.net/sinat_22336563/article/details/68496919

    相关文章

      网友评论

        本文标题:Caffe源码解读:syncedmem类

        本文链接:https://www.haomeiwen.com/subject/chybkxtx.html