caffe源码阅读——SyncedMemory.cpp

最新推荐文章于 2020-02-28 13:37:00 发布

kuguadaoren

最新推荐文章于 2020-02-28 13:37:00 发布

阅读量199

点赞数

分类专栏： caffe 文章标签：源码阅读

本文链接：https://blog.csdn.net/kuguadaoren/article/details/78166927

版权

caffe 专栏收录该内容

2 篇文章 0 订阅

订阅专栏

SyncedMemory的主要功能是实现CPU和GPU中的数据进行同步。

SyncedMemory的构造函数

//SyncedMemory的构造函数
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
}

SyncedMemory的析构函数

SyncedMemory::~SyncedMemory() {
  check_device();
  if (cpu_ptr_ && own_cpu_data_) {
    CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);//释放CPU数据内存
  }

#ifndef CPU_ONLY
  if (gpu_ptr_ && own_gpu_data_) {
    CUDA_CHECK(cudaFree(gpu_ptr_));//释放GPU数据内存
  }
#endif  // CPU_ONLY
}

SyncedMemory::to_cpu()

inline void SyncedMemory::to_cpu() {    //实现把数据存入cpu
  check_device();
  switch (head_) {
  case UNINITIALIZED:    //数据未初始化，则在cpu申请内存
    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:     //数据在GPU上存储时，把数据拷贝到CPU
#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:   //数据存储在CPU时，不做任何处理
  case SYNCED:        //数据在CPU和gpu上时不做任何处理
    break;
  }
}

SyncedMemory::to_gpu()

和SyncedMemory::to_cpu()类似，主要实现在GPU上的数据存储。

SyncedMemory::cpu_data()

const void* SyncedMemory::cpu_data() {  //返回指向存储cpu数据的指针cpu_ptr_
  check_device();
  to_cpu();
  return (const void*)cpu_ptr_;
}

SyncedMemory::set_cpu_data(void* data)

void SyncedMemory::set_cpu_data(void* data) {//指向cpu内存的指针cpu_ptr_指向新的地方，并释放原先申请的内存
  check_device();
  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;
}

SyncedMemory::gpu_data()

和SyncedMemory::cpu_data()类似，主要功能为返回指向存储gpu数据的指针gpu_ptr_。

SyncedMemory::set_gpu_data(void* data)

和SyncedMemory::set_cpu_data(void* data)类似，主要功能释放GPU存储的数据。

SyncedMemory::mutable_cpu_data()

void* SyncedMemory::mutable_cpu_data() {//返回指向CPU数据的指针，并把head_的状态设置为HEAD_AT_CPU
  check_device();
  to_cpu();
  head_ = HEAD_AT_CPU;
  return cpu_ptr_;
}

SyncedMemory::mutable_gpu_data()

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
}//返回指向GPU数据的指针，并把head_的状态设置为HEAD_AT_GPU