Caffe --- SyncedMemory

SyncedMemory类定义在syncedmem.hpp/cpp里, 主要负责caffe底层的内存管理.

PS: Caffe的底层数据的切换（cpu模式和gpu模式），需要用到内存同步模块。其实个人觉得如果需要研究Blob，对于SyncedMemory的分析很重要

内存分配与释放

内存分配与释放由两个(不属于SyncedMemory类的)内联函数完成.
代码简单直观:

如果是CPU模式, 那么调用malloc和free来申请/释放内存,

否则调用CUDA的cudaMallocHost和cudaFreeHost来申请/释放显存.

// ------ 分配内存 ------ 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";}// ------ 释放内存 ------ 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 内存地址  void* gpu_ptr_;  // gpu 内存地址  size_t size_;  // 数据大小  SyncedHead head_;  // 当前数据同步状态  bool own_cpu_data_;  //  是否是自己的cpu data? (例如set_cpu_data就是false)  bool cpu_malloc_use_cuda_;  bool own_gpu_data_;  // 是否已经申请gpu内存空间  int gpu_device_;  //

get and set 方法

cpu_data, gpu_data或者mutable_cpu_data, mutable_gpu_data方法返回cpu或者gpu内存指针, 前者是const void*, 不可对返回内存进行修改; 后者为void*, 可以修改.

set方法比较特别, 方法参数是指向另一段内存空间的地址:

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;}

该函数首先释放自己申请的内存空间, 然后直接指向参数传入的内存空间 (并不是重新申请空间, 并copy数据). 最后将 own_cpu_data_设置为false, 表示外来数据(?).

保持数据同步

在调用cpu_data或者gpu_data方法时, 需要确保cpu, gpu数据内容是一致的. 这里用到了前面提到的枚举类型来记录当前同步状态

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

以to_cpu()方法为例: 检查head_所处状态, 若UNINITIALIZED, 则分配内存空间(置0); 若HEAD_AT_GPU, 则需要从GPU内存同步数据到CPU; HEAD_AT_CPU, 则说明目前最新的数据是在CPU的, 无须进行任何操作 (虽然并不知道GPU的数据是否和CPU一致, 因为当前我们并不关心GPU数据); 若SYNCED, 则CPU/GPU数据一致, 无须进行任何操作.

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;  }}

总结

通过读SyncedMemory的源代码，可以看到，用它可以为数据分配内存空间，根据需要储存，获取，修改cpu或gpu数据，比如，set_cpu_data，set_gpu_data，cpu_data，gpu_data，mutable_cpu_data，mutable_gpu_data。并可以在cpu和gpu之间进行同步，不用关心其具体细节。

该部分引用博主牛闯的总结。

源码

syncedmem.hpp

#ifndef CAFFE_SYNCEDMEM_HPP_#define CAFFE_SYNCEDMEM_HPP_#include <cstdlib>#include "caffe/common.hpp"namespace caffe {// If CUDA is available and in GPU mode, host memory will be allocated pinned,// using cudaMallocHost. It avoids dynamic pinning for transfers (DMA).// The improvement in performance seems negligible in the single GPU case,// but might be more significant for parallel training. Most importantly,// it improved stability for large models on many GPUs.// 如果使用cuda，那么用cudaMallocHost分配管理主机内存，这种方式在多gpu上并行计算时，性能会显著提高。// 如果只用cpu，则用malloc分配内存。要想了解两者的差异，还需进一步学习。// 分配内存函数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";}// 释放内存函数inline void CaffeFreeHost(void* ptr, bool use_cuda) {#ifndef CPU_ONLY  if (use_cuda) {    CUDA_CHECK(cudaFreeHost(ptr));    return;  }#endif  free(ptr);}/** * @brief Manages memory allocation and synchronization between the host (CPU) *        and device (GPU). * * TODO(dox): more thorough description. */ // 管理内存分配和同步的类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数据的void类型的指针，用void类型可以管理任意类型的数据，对于具体类型数据只需制定指针类型即可  // 另外指针时const的，不能通过该指针改变数据  const void* cpu_data();  // 将当前cpu_ptr_指向data指向的数据，并将其原来指向的数据（如果存在）释放  void set_cpu_data(void* data);  // 下面两个与cpu类似  const void* gpu_data();  void set_gpu_data(void* data);  // 下面两个也与cpu_data类似，区别是可以通过该指针改变其数据  void* mutable_cpu_data();  void* mutable_gpu_data();  // head的状态，前三个分别是，没有初始化，在cpu，在gpu，最后一个表示同步了，说名数据刚从cpu转到gpu，或gpu到cpu  // 下面的函数要很据这些状态来判断是否同步，怎样同步  enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };  SyncedHead head() { return head_; }  // 返回size  size_t size() { return size_; }#ifndef CPU_ONLY  // 将cpu的数据同步到gpu上，并head_ = SYNCED  void async_gpu_push(const cudaStream_t& stream);#endif private:  // 将数据同步到cpu，根据head的状态进行不同的操作，具体的可以看.cpp文件  void to_cpu();  // 与cpu类似  void to_gpu();  // 下面就是其私有成员变量了，见名知意  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_;  // 这个其实就是禁止使用复制和赋值操作符（=）  DISABLE_COPY_AND_ASSIGN(SyncedMemory);};  // class SyncedMemory}  // namespace caffe#endif  // CAFFE_SYNCEDMEM_HPP_

syncedmem.cpp

#include "caffe/common.hpp"#include "caffe/syncedmem.hpp"#include "caffe/util/math_functions.hpp"namespace caffe {SyncedMemory::~SyncedMemory() {  if (cpu_ptr_ && own_cpu_data_) {    CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);  }#ifndef CPU_ONLY  if (gpu_ptr_ && own_gpu_data_) {    int initial_device;    cudaGetDevice(&initial_device);    if (gpu_device_ != -1) {      CUDA_CHECK(cudaSetDevice(gpu_device_));    }    CUDA_CHECK(cudaFree(gpu_ptr_));    cudaSetDevice(initial_device);  }#endif  // CPU_ONLY}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;  }}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_);    head_ = SYNCED;    break;  case HEAD_AT_GPU:  case SYNCED:    break;  }#else  NO_GPU;#endif}const void* SyncedMemory::cpu_data() {  to_cpu();  return (const void*)cpu_ptr_;}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;}const void* SyncedMemory::gpu_data() {#ifndef CPU_ONLY  to_gpu();  return (const void*)gpu_ptr_;#else  NO_GPU;  return NULL;#endif}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_));    }    CUDA_CHECK(cudaFree(gpu_ptr_));    cudaSetDevice(initial_device);  }  gpu_ptr_ = data;  head_ = HEAD_AT_GPU;  own_gpu_data_ = false;#else  NO_GPU;#endif}void* SyncedMemory::mutable_cpu_data() {  to_cpu();  head_ = HEAD_AT_CPU;  return cpu_ptr_;}void* SyncedMemory::mutable_gpu_data() {#ifndef CPU_ONLY  to_gpu();  head_ = HEAD_AT_GPU;  return gpu_ptr_;#else  NO_GPU;  return NULL;#endif}#ifndef CPU_ONLYvoid 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

Reference

http://www.cnblogs.com/louyihang-loves-baiyan/p/5150554.html
platero（简书作者）