Hacks on Computer Vision

caffe syncedmem 源码简析

2016.04.19

最近想把caffe简化一下,于是又开始了看源码。刚看完了syncedmem的代码,写个总结记录一下~

从名字看,syncedmem是内存同步相关的功能,即数据从内存到GPU显存的同步。

syncedmem.hpp文件中,有下面这两个内联函数,简单的对mallocfree内存做了封装:

inline void CaffeMallocHost(void** ptr, size_t size) {
  *ptr = malloc(size);
}

inline void CaffeFreeHost(void* ptr) {
  free(ptr);
}

从注释上看,作者原想也封装在GPU的内存,但在使用 cuda 5.0 时遇到些问题,就简单这样处理。

下面就是类SyncedMemory的声明:

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();
  void set_cpu_data(void* data);
  const void* gpu_data();
  void* mutable_cpu_data();
  void* mutable_gpu_data();
  enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
  SyncedHead head() { return head_; }
  size_t size() { return size_; }

 private:
  void to_cpu();
  void to_gpu();
  void* cpu_ptr_;
  void* gpu_ptr_;
  size_t size_;
  SyncedHead head_;
  bool own_cpu_data_;

  DISABLE_COPY_AND_ASSIGN(SyncedMemory);
};  // class SyncedMemory

这里首先要看的是private里面的这五个私有变量:

void* cpu_ptr_;
void* gpu_ptr_;
size_t size_;
SyncedHead head_;
bool own_cpu_data_;

因为构建函数会对它们做初始化,所以了解它们的含义比较重要。其中,除了SyncedHead head_之外,其他四个的含义都比较直接,前两个分别是cpugpu上的数据指针,第三个是数据的大小,最后一个bool类型表示是否存在cpu数据。

关于SyncedHead head_,前面是有定义的:

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

这四个分别表明了数据的状态:未初始化、在CPU、在GPU、已同步。这个在具体实现的代码syncedmem.cpp中可以看出来,比如这个函数to_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;
  }
}

这个函数的目的是把数据传到cpu上,那如果数据状态是未初始化,就需要创建数据,并把head_设为HEAD_AT_CPU,同时把own_cpu_data_设为true;而如果数据在gpu上,则把数据从gpu复制到cpu的数据中,并把head_设为SYNCED,同时把own_cpu_data_设为true

通过syncedmem的控制,可以把训练过程中的数据在cpu和gpu之间同步,另外因为Blob中使用的是shared_ptr,二者结合,内存上也是一种优化的方式。

注:我所使用的caffe版本较早,与最新版可能存在差异。

__EOF__

本文作者HackCV
版权声明本博客所有文章除特别声明外,均采用 BY-NC-SA 许可协议。转载请注明出处!
本文链接https://hackcv.com/posts/caffe-syncedmem-%E6%BA%90%E7%A0%81%E7%AE%80%E6%9E%90/

发表评论