caffe源码 SyncedMemory 详解

最新推荐文章于 2019-11-09 21:14:14 发布

置顶兢兢业业小码农

最新推荐文章于 2019-11-09 21:14:14 发布

阅读量678

点赞数 1

CC 4.0 BY-SA版权

分类专栏： caffe源码 caffe | pytorch 源码共读项目分享文章标签： caffe源码 syncedmemory

本文链接：https://blog.youkuaiyun.com/qq_28660035/article/details/80347208

caffe | pytorch 源码共读项目分享同时被 2 个专栏收录

17 篇文章

订阅专栏

caffe源码

13 篇文章

订阅专栏

本文详细解析了Caffe中的SyncedMemory类，该类负责管理内存分配和GPU、CPU间的数据交互。文章介绍了SyncedMemory的不同状态及其转换逻辑，包括如何在CPU和GPU之间同步数据，以及如何设置和获取数据。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

SyncedMemory是caffe中直接用于管理内存分配和GPU、CPU数据交互的类，只要服务Blob类。

其头文件中大部分是声明，还有一些特别的：

优化多GPU：如果用CUDA+GPU模式，CPU内存将会被分配固定(分配的内存可以常驻在内存空间中对效率是有帮助的，空间不会被别的进程所抢占)用cudamallochost，避免了用于传输的动态固定。这种对单GPU没啥用，但对多GPU提升了大模型的稳定性和并行训练的效率。

构造函数输入类型：

explicit SyncedMemory(size_t size);//typedef __SIZE_TYPE__ size_t,其中__SIZE_TYPE__为long unsigned int

四种状态描述：

//四种状态量，表示数据处于不同状态：未初始化，在CPU，在GPU，GPU和CPU同步，不同状态调用函数操作也不同：

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

SyncedHead head() { return head_; }

SyncedHead head_;

举个栗子，如果调用to_gpu(),不同状态反映不同，如果是HEAD_AT_GPU, SYNCED则不进行任何操作，因为已经在gpu上了，

但如果是HEAD_AT_CPU,则将数据拷到从cpu同步到gpu，并将own_gpu_data设为true，表示该blob拥有自己开辟可释放的gpu上数据，并将haed_设为SYNCED；如果head_是未被初始化的状态UNINITIALIZED，那么首先需要先分配内存，之后再将开辟空间内数据都置0，并将own_gpu_data设为true，表示该blob拥有自己开辟可释放的gpu上数据，并将haed_设为HEAD_AT_GPU。

详情可见（https://blog.youkuaiyun.com/u010414386/article/details/52346192）

私有成员：
void check_device();
void to_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 device_;

下面是syncedmen.cpp源码及注释：

#include "caffe/common.hpp"
#include "caffe/syncedmem.hpp"
#include "caffe/util/math_functions.hpp"

namespace caffe {
//构造函数，显式初始化参数
SyncedMemory::SyncedMemory()
  : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
    own_cpu_data_(false), cpu_malloc_use_cuda_(false), own_gpu_data_(false) {
#ifndef CPU_ONLY
#ifdef DEBUG
  CUDA_CHECK(cudaGetDevice(&device_));
#endif
#endif
}

SyncedMemory::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) {
#ifndef CPU_ONLY
#ifdef DEBUG
  CUDA_CHECK(cudaGetDevice(&device_));
#endif
#endif
}
//析构函数，释放CPU GPU上的内存空间
SyncedMemory::~SyncedMemory() {
  check_device();
  if (cpu_ptr_ && own_cpu_data_) {
    CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
  }

#ifndef CPU_ONLY
  if (gpu_ptr_ && own_gpu_data_) {
    CUDA_CHECK(cudaFree(gpu_ptr_));
  }
#endif  // CPU_ONLY
}
//根据不同状态，将数据转移到cpu。返回const void*, 不可对返回内存进行修改
inline void SyncedMemory::to_cpu() {
  check_device();
  switch (head_) {
  case UNINITIALIZED:
    CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);//分配内存
    caffe_memset(size_, 0, cpu_ptr_);//将分配的内存全部初始化为0
    head_ = HEAD_AT_CPU;
    own_cpu_data_ = true;//表示现在有通过这种方式分配的cpu空间，如果要设置指向新数据，要先释放现有数据内存空间。
    break;
  //若在GPU上，先在CPU分配空间，然后再利用指针同步
  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;
  }
}
//根据不同状态，将数据转移到gpu。返回const void*, 不可对返回内存进行修改
inline void SyncedMemory::to_gpu() {
  check_device();
#ifndef CPU_ONLY
  switch (head_) {
  case UNINITIALIZED:
    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(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() {
  check_device();
  to_cpu();
  return (const void*)cpu_ptr_;
}
//返回void*，可以修改，只是指向另一段内存空间地质
void SyncedMemory::set_cpu_data(void* data) {
  check_device();
  CHECK(data);
  //判断现在是否有cpu上数据，如果有，要先释放现拥有的数据内存空间。
  if (own_cpu_data_) {
    CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
  }
  cpu_ptr_ = data;//指向另一段内存空间地址，并不是重新申请空间, 并copy数据
  head_ = HEAD_AT_CPU;
  own_cpu_data_ = false;
}

const void* SyncedMemory::gpu_data() {
  check_device();
#ifndef CPU_ONLY
  to_gpu();
  return (const void*)gpu_ptr_;
#else
  NO_GPU;
  return NULL;
#endif
}

void SyncedMemory::set_gpu_data(void* data) {
  check_device();
#ifndef CPU_ONLY
  CHECK(data);
  if (own_gpu_data_) {
    CUDA_CHECK(cudaFree(gpu_ptr_));
  }
  gpu_ptr_ = data;//指向另一段内存空间地址，并不是重新申请空间, 并copy数据
  head_ = HEAD_AT_GPU;
  own_gpu_data_ = false;
#else
  NO_GPU;
#endif
}
//获取cpu上数据，并返回可修改的void*类型指针
void* SyncedMemory::mutable_cpu_data() {
  check_device();
  to_cpu();
  head_ = HEAD_AT_CPU;
  return cpu_ptr_;
}
//获取gpu上数据，并返回可修改的void*类型指针
void* SyncedMemory::mutable_gpu_data() {
  check_device();
#ifndef CPU_ONLY
  to_gpu();
  head_ = HEAD_AT_GPU;
  return gpu_ptr_;
#else
  NO_GPU;
  return NULL;
#endif
}
//异步同步数据流，就是实现cpu的数据复制到gpu里面。使用异步方式可以有效的防止主进程阻塞。
#ifndef CPU_ONLY
void SyncedMemory::async_gpu_push(const cudaStream_t& stream) {
  check_device();
  CHECK(head_ == HEAD_AT_CPU);
  if (gpu_ptr_ == NULL) {
    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

void SyncedMemory::check_device() {
#ifndef CPU_ONLY
#ifdef DEBUG
  int device;
  cudaGetDevice(&device);
  CHECK(device == device_);
  if (gpu_ptr_ && own_gpu_data_) {
    cudaPointerAttributes attributes;
    CUDA_CHECK(cudaPointerGetAttributes(&attributes, gpu_ptr_));
    CHECK(attributes.device == device_);
  }
#endif
#endif
}

}  // namespace caffe

注意：

inline void SyncedMemory::to_cpu/gpu()，返回const void*, 不可对返回内存进行修改。
void* SyncedMemory::mutable_cpu/gpu_data()，返回可修改的void*类型指针
void SyncedMemory::set_cpu/gpu_data，返回void*，可以修改，指向另一段内存空间地址

我们可以看到set_cpu_data释放了当前的cpu内存，把指针指向data所指的内存中，own_cpu_data_ 设置为了false；表明当前使用的是宿主(data)的内存。

我们对own_cpu_data_ 进行标记是有必要的，因为当使用的是宿主的内存的时候，当这个类被释放而调用析构函数时，需要检查共享标记，不能释放宿主的内存。

这样可以保证自己申请的内存只能由自己释放。