caffe源码分析--SyncedMemory 内存管理机制

NH558J0

于 2019-11-09 21:14:14 发布

阅读量430

点赞数

分类专栏： caffe 文章标签： caffe源码分析 syncedmemory

本文链接：https://blog.csdn.net/happyday_d/article/details/102991885

版权

caffe 专栏收录该内容

5 篇文章 0 订阅

订阅专栏

caffe源码分析–SyncedMemory 内存管理机制

SyncedMemory 是caffe中用来管理内存分配和CPU、GPU数据及同步的类，只服务于Blob类。SyncedMemory 对象管理的是一个tensor的数据对象，这个对象可能只存在CPU上，也有可能存在GPU上，或者同时在两个位置上，如何保证在CPU和GPU上的数据能同步呢？

1. 自动机模型管理

SyncedMemory 使用自动机模型对内存和显存进行管理。对每个SyncedMemory 对象都包含了一个*head_*的枚举变量。*head_*变量是SyncedMemory 对象内部变量，标志着目前被管理的数据对象的同步情况，定义如下：

private:
  void to_cpu(); //数据由显存同步到内存
  void to_gpu(); //数据由内存同步到显存
  void* cpu_ptr_;//内存指针
  void* gpu_ptr_;//显存指针
  size_t size_;  //数据大小
  SyncedHead head_;//当前数据状态，UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED

  //own_cpu_data_和own_gpu_data_这两个变量。这两个变量主要是用来记录是否使用了共享的数据还是自己的数据
  // 这里own_cpu_data_和own_gpu_data_不是互斥的关系，可以同时own两个地方的数据
  // 有一种可能是SyncedMemory对象所包含的数据指针，是指向的另外一段内存空间，而不由自己申请
  bool own_cpu_data_;
  bool cpu_malloc_use_cuda_;
  bool own_gpu_data_;
  int gpu_device_;

其中：*cpu_ptr_*指向CPU上的数据单元，gpu_ptr_*指向GPU上的数据单元。

SyncedHead定义如下：

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

SyncedHead是一个枚举类型，所以*head_*可取的值为：UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED

UNINITIALIZED:表示这个数据对象的存储空间没有被初始化，即SyncedMemory 最早的状态，内存和显存都没有被分配，当CPU或者GPU申请内存时该状态终结。
HEAD_AT_CPU：表明最近一次数据修改是由CPU引起的，CPU上的数据时最新的，此时CPU和GPU的数据还没有同步，CPU和GPU数据可能不同。
HEAD_AT_GPU：表明最近一次数据修改是由GPU引起的，GPU上的数据时最新的，此时CPU和GPU的数据还没有同步，CPU和GPU数据可能不同。
SYNCED：同步状态，表明CPU和GPU的数据一致。

自动机的四个状态是如何转换呢？

参考下图：

在这里插入图片描述

他们会被相应的状态转移函数触发: to_cpu()、to_gpu()、mutable_cpu_data()、mutable_gpu_data()。其中mutable_cpu_data()、mutable_gpu_data()函数主要目的是为了得到可修改的cpu和gpu指针，具体实现如下：

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
}

mutable_cpu_data()调用to_cpu()函数、mutable_gpu_data()调用to_gpu()函数，并将*head_*修改为相应的值。

to_cpu()函数实现如下：

//函数作用是让现在数据最新备份至少出现在cpu上
//函数执行过程中，如果最新的数据在GPU上，那么把GPU的数据同步到CPU上
//这时候CPU和GPU都有最新的同步数据
//返回状态机结果是HEAD_AT_CPU或者SYNCED
//HEAD_AT_CPU -- 代表数据备份最新出现在CPU上, GPU上的数据不是最新的备份
//SYNCED -- 代表CPU和GPU数据都是最新
inline void SyncedMemory::to_cpu() {
  switch (head_) {
  case UNINITIALIZED:
  //如果SyncedMemory对象所对应的这段内存没有被初始化，那么在host上申请内存
    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上
#ifndef CPU_ONLY
// 首先在Host上申请这么一段内存，然后把GPU上的数据，同步到CPU上，然后把head状态更新为同步
    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;
  }
}

own_cpu_data_ 的作用是什么？

如前注释：*own_cpu_data_和own_gpu_data_这两个变量，这两个变量主要是用来记录是否使用了共享的数据还是自己的数据，to_cpu()代码中为cpu分配了内存，此时own_cpu_data_*设置为true,表示使用了当前自己的数据。而在set_cpu_date函数中：

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

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

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

to_gpu()函数实现如下：

//函数作用是让现在数据最新备份至少出现在GPU上
//函数执行过程中，如果最新的数据在CPU上，那么把CPU的数据同步到GPU上
//这时候CPU和GPU都有最新的同步数据
//返回状态机结果是HEAD_AT_CPU或者SYNCED
//HEAD_AT_GPU -- 代表数据备份最新出现在CPU上, CPU上的数据不是最新的备份
//SYNCED -- 代表CPU和GPU数据都是最新的
inline void SyncedMemory::to_gpu() {
#ifndef CPU_ONLY
  switch (head_) {
    //如果指针数据没有被分配，那么在GPU上重新分配
  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;
    //如果目前数据在CPU head上，那么把CPU上的数据，拷贝到GPU上
  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
}

什么时候改变状态机变量的值，即什么时候调用to_cpu, to_gpu函数

在blob对象访问CPU或者GPU数据指针的时候。

参考下面的调用关系：

在conv_layer.cpp中Forward_cpu函数调用的是blob的cpu_data()函数，这个函数返回的是数据的实际指针。

template <typename Dtype>
void ConvolutionLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {
  const Dtype* weight = this->blobs_[0]->cpu_data();

  for (int i = 0; i < bottom.size(); ++i) {
    const Dtype* bottom_data = bottom[i]->cpu_data();
    Dtype* top_data = top[i]->mutable_cpu_data();

在blob的cpu_data()函数如下：

template <typename Dtype>
const Dtype* Blob<Dtype>::cpu_data() const {
  CHECK(data_);
  return (const Dtype*)data_->cpu_data();   // 调用SyncedMemory的数据访问函数cpu_data(),并返回内存指针
}

SyncedMemory对象的cpu_data()函数定义如下：

const void* SyncedMemory::cpu_data() {
  to_cpu();                 // 首先完成数据同步，第一次访问时会申请存储空间
  return (const void*)cpu_ptr_;
}

blob中*data_ *的定义如下：

 protected:
  shared_ptr<SyncedMemory> data_; //存储前向传递数据
  shared_ptr<SyncedMemory> diff_; //存储反向传递梯度
  shared_ptr<SyncedMemory> shape_data_;
  vector<int> shape_;  //参数维度
  int count_; //Blob存储的元素个数（shape_所有元素乘积）
  int capacity_;//当前Blob的元素个数（控制动态分配）

2.源码注释

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使用memory pinned技术,使用cudaMallocHost分配的内存，这样可以加速CPU和GPU数据的传输速度
//在多GPU训练下能起到大的作用
inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
#ifndef CPU_ONLY
  if (Caffe::mode() == Caffe::GPU) {
    CUDA_CHECK(cudaMallocHost(ptr, size));//GPU下使用cuda分配内存
    *use_cuda = true;
    return;
  }
#endif
  *ptr = malloc(size); //如果只是用cpu则用c的malloc分配内存
  *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(); //析构函数，调用CaffeFreeHost释放内存
  const void* cpu_data();        //获得cpu数据指针
  void set_cpu_data(void* data); //设置cpu使用共享数据
  const void* gpu_data();        //获得gpu数据指针
  void set_gpu_data(void* data); //设置gpu使用共享数据
  void* mutable_cpu_data();      //获取可更改的cpu数据指针
  void* mutable_gpu_data();      //获取可更改的gpu数据指针
  enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED }; //自动状态转换机状态
  SyncedHead head() { return head_; } //获取当前数据状态
  size_t size() { return size_; }

#ifndef CPU_ONLY
  void async_gpu_push(const cudaStream_t& stream);
#endif

 private:
  void to_cpu(); //数据由显存同步到内存
  void to_gpu(); //数据由内存同步到显存
  void* cpu_ptr_;//内存指针
  void* gpu_ptr_;//显存指针
  size_t size_;  //数据大小
  SyncedHead head_;//当前数据状态，UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED

  // 这里own_cpu_data_和own_gpu_data_不是互斥的关系，可以同时own两个地方的数据
  // own_cpu_data_和own_gpu_data_这两个变量。这两个变量主要是用来记录是否使用了共享的数据还是自己的数据
  // 有一种可能是SyncedMemory对象所包含的数据指针，是指向的另外一段内存空间，而不由自己申请
  bool own_cpu_data_; //共享标记，是否使用的是自己的cpu数据
  bool cpu_malloc_use_cuda_;
  bool own_gpu_data_; //共享标记，是否使用的是自己的gpu数据
  int gpu_device_;    //gpu设备

  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_);  //own_cpu_data_为true时释放cpu内存，保证不释放宿主数据
  }

#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_)); //释放gpu内存
    cudaSetDevice(initial_device);
  }
#endif  // CPU_ONLY
}
//函数作用是让现在数据最新备份至少出现在cpu上
//函数执行过程中，如果最新的数据在GPU上，那么把GPU的数据同步到CPU上
//这时候CPU和GPU都有最新的同步数据
//返回状态机结果是HEAD_AT_CPU或者SYNCED
//HEAD_AT_CPU -- 代表数据备份最新出现在CPU上, GPU上的数据不是最新的备份
//SYNCED -- 代表CPU和GPU数据都是最新
inline void SyncedMemory::to_cpu() {
  switch (head_) {
  case UNINITIALIZED:
  //如果SyncedMemory对象所对应的这段内存没有被初始化，那么在host上申请内存
    CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
    caffe_memset(size_, 0, cpu_ptr_); //将分配的内存全部初始化为0
    head_ = HEAD_AT_CPU;
    own_cpu_data_ = true;
    break;
  case HEAD_AT_GPU: //表示目前数据是在GPU上
#ifndef CPU_ONLY
// 首先在Host上申请这么一段内存，然后把GPU上的数据，同步到CPU上，然后把head状态更新为同步
    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上
//函数执行过程中，如果最新的数据在CPU上，那么把CPU的数据同步到GPU上
//这时候CPU和GPU都有最新的同步数据
//返回状态机结果是HEAD_AT_CPU或者SYNCED
//HEAD_AT_GPU -- 代表数据备份最新出现在CPU上, CPU上的数据不是最新的备份
//SYNCED -- 代表CPU和GPU数据都是最新的
inline void SyncedMemory::to_gpu() {
#ifndef CPU_ONLY
  switch (head_) {
    //如果指针数据没有被分配，那么在GPU上重新分配
  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;
    //如果目前数据在CPU head上，那么把CPU上的数据，拷贝到GPU上
  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
}
//获得cpu数据指针
const void* SyncedMemory::cpu_data() {
  to_cpu();                 // 首先完成数据同步，第一次访问时会申请存储空间
  return (const void*)cpu_ptr_;
}
//设置cpu共享数据
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;
}
//获得gpu数据指针
const void* SyncedMemory::gpu_data() {
#ifndef CPU_ONLY
  to_gpu();
  return (const void*)gpu_ptr_;
#else
  NO_GPU;
  return NULL;
#endif
}
//设置gpu共享数据
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
}

//async_gpu_push 这个函数的作用是异步同步数据流，就是实现cpu的数据复制到gpu里面。使用异步方式可以有效的防止主进程阻塞。
#ifndef CPU_ONLY
void 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

参考：

1.https://blog.csdn.net/qq_28660035/article/details/80347208

2.https://blog.csdn.net/hnshahao/article/details/81218713

NH558J0

关注

0
点赞
踩
1

收藏

觉得还不错? 一键收藏
0
评论
caffe源码分析--SyncedMemory 内存管理机制

caffe源码分析–SyncedMemory 内存管理机制 SyncedMemory 是caffe中用来管理内存分配和CPU、GPU数据及同步的类，只服务于Blob类。SyncedMemory 对象管理的是一个tensor的数据对象，这个对象可能只存在CPU上，也有可能存在GPU上，或者同时在两个位置上，如何保证在CPU和GPU上的数据能同步呢？1. 自动机模型管理 SyncedMemo...
复制链接

扫一扫