Caffe源码解读:syncedmem类
2017-07-24 本文已影响103人
Mordekaiser
内存同步(syncedmem)类的作用在于管理主机(CPU)和设备(GPU)之间的内存分配和数据同步,封装了二者之间的交互操作。
这个类没有对应的ProtoBuffer描述,所以直接看./include/caffe/syncedmem.cpp文件:
#ifndef CAFFE_SYNCEDMEM_HPP_
#define CAFFE_SYNCEDMEM_HPP_
#include <cstdlib>
#include "caffe/common.hpp"
namespace caffe {
// 以页锁定方式分配内存,在单GPU时提示不明显,多GPU提升很多
// malloc分配标准可分页的主机内存,操作系统可能会将这种内存分页或者交换到磁盘上,需要的时候调回内存,这样可能会增加运行时间
// cudaMallocHost分配页锁定的主机内存,操作系统不会对这块内存分页或者交换到磁盘上,可以节省时间
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); // CPU模式下分配内存
*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) {
//如果分配内存用的是cudaMallocHost分配,即use_cuda为真,cpu中的数据也可以用cudaMallocHost分配内存
CUDA_CHECK(cudaFreeHost(ptr));
return;
}
#endif
free(ptr); // 用的malloc分配的内存
}
// 负责内存分配和设备同步
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,GPU数据的读写,不赘述。为什么没有set_cpu_diff() ???
const void* cpu_data();
void set_cpu_data(void* data);
const void* gpu_data();
void set_gpu_data(void* data);
void* mutable_cpu_data();
void* mutable_gpu_data();
// 共享内存的4种状态:未初始化,CPU数据有效,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(); // 数据同步至cpu
void to_gpu(); // 数据同步至gpu
void* cpu_ptr_; // cpu中数据的指针
void* gpu_ptr_; // gpu中数据的指针
size_t size_; // 存储空间的大小
SyncedHead head_; // 共享内存的状态
bool own_cpu_data_; // cpu拥有数据所有权
bool cpu_malloc_use_cuda_; // 分配cpu内存是否用cudaMallocHost()分配。
bool own_gpu_data_; // gpu拥有数据所有权
int gpu_device_; // gpu设备号
// 禁用拷贝构造以及赋值运算符
// 使用grep可以查到,该宏定义在common.hpp第35行
// 该宏就是把拷贝构造和赋值运算符设置为private而已
DISABLE_COPY_AND_ASSIGN(SyncedMemory);
}; // class SyncedMemory
} // namespace caffe
#endif // CAFFE_SYNCEDMEM_HPP_
这个类比Blob简单多了,下面看对应的./src/caffe/syncedmem.cpp文件:
#include "caffe/common.hpp"
#include "caffe/syncedmem.hpp"
#include "caffe/util/math_functions.hpp"
namespace caffe {
SyncedMemory::~SyncedMemory() {
// 如果cpu拥有数据所有权
if (cpu_ptr_ && own_cpu_data_) {
CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
}
// 如果数据在gpu上
#ifndef CPU_ONLY
if (gpu_ptr_ && own_gpu_data_) {
int initial_device;
cudaGetDevice(&initial_device); // 获取使用的gpu设备号
if (gpu_device_ != -1) {
CUDA_CHECK(cudaSetDevice(gpu_device_));
}
CUDA_CHECK(cudaFree(gpu_ptr_));
cudaSetDevice(initial_device);
}
#endif // CPU_ONLY
}
// 数据同步到cpu上
inline void SyncedMemory::to_cpu() {
switch (head_) {
case UNINITIALIZED: // 如果是未初始化状态,只需分配下内存就行了
CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
caffe_memset(size_, 0, cpu_ptr_); // 定义在math_functions.hpp第42行,调用了memset
head_ = HEAD_AT_CPU; // 内存状态改为cpu拥有所有权
own_cpu_data_ = true;
break;
case HEAD_AT_GPU: // 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;
// 数据已经为cpu拥有所有权或者在内存共享状态,则什么都不管
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_); //数据拷贝,调用了cudaMemcpy函数
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_) {
// 调用这个函数的时候,如果cpu内有数据会被直接清空,要注意
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_));
}
// 调用这个函数的时候,如果gpu内有数据会被直接清空,要注意
CUDA_CHECK(cudaFree(gpu_ptr_));
cudaSetDevice(initial_device);
}
gpu_ptr_ = data;
head_ = HEAD_AT_GPU;
own_gpu_data_ = false;
#else
NO_GPU;
#endif
}
// 读写获取cpu数据
void* SyncedMemory::mutable_cpu_data() {
to_cpu();
head_ = HEAD_AT_CPU;
return cpu_ptr_;
}
// 读写获取gpu数据
void* SyncedMemory::mutable_gpu_data() {
#ifndef CPU_ONLY
to_gpu();
head_ = HEAD_AT_GPU;
return gpu_ptr_;
#else
NO_GPU;
return NULL;
#endif
}
// cuda中的流同步,这里传入一个异步流,在计算的时候向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
syncedmem类比较简单,主要是完成cpu和gpu之间的数据交互问题~
参考资料
- 《21天实战caffe》
- (介绍了流同步)http://blog.csdn.net/sinat_22336563/article/details/68496919