caffe syncedmem 源码简析

2016.04.19

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

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

在syncedmem.hpp文件中，有下面这两个内联函数，简单的对malloc和free内存做了封装：

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_之外，其他四个的含义都比较直接，前两个分别是cpu和gpu上的数据指针，第三个是数据的大小，最后一个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/

caffe 源码阅读 syncedmem

发表评论

HackCV

caffe syncedmem 源码简析

相关文章