最近想把caffe
简化一下,于是又开始了看源码。刚看完了syncedmem
的代码,写个总结记录一下~
从名字看,syncedmem
是内存同步相关的功能,即数据从内存到GPU显存的同步。
在syncedmem.hpp
文件中,有下面这两个内联函数,简单的对malloc
和free
内存做了封装:
inline void CaffeMallocHost(void** ptr, size_t size) {
*ptr = malloc(size);
}
inline void CaffeFreeHost(void* ptr) {
free(ptr);
}
从注释上看,作者原想也封装在GPU的内存,但在使用 cuda 5.0 时遇到些问题,就简单这样处理。
下面就是类SyncedMemory
的声明:
class SyncedMemory {
public:
SyncedMemory()
: cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
own_cpu_data_(false) {}
explicit SyncedMemory(size_t size)
: cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED),
own_cpu_data_(false) {}
~SyncedMemory();
const void* cpu_data();
void set_cpu_data(void* data);
const void* gpu_data();
void* mutable_cpu_data();
void* mutable_gpu_data();
enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
SyncedHead head() { return head_; }
size_t size() { return size_; }
private:
void to_cpu();
void to_gpu();
void* cpu_ptr_;
void* gpu_ptr_;
size_t size_;
SyncedHead head_;
bool own_cpu_data_;
DISABLE_COPY_AND_ASSIGN(SyncedMemory);
}; // class SyncedMemory
这里首先要看的是private
里面的这五个私有变量:
void* cpu_ptr_;
void* gpu_ptr_;
size_t size_;
SyncedHead head_;
bool own_cpu_data_;
因为构建函数会对它们做初始化,所以了解它们的含义比较重要。其中,除了SyncedHead head_
之外,其他四个的含义都比较直接,前两个分别是cpu
和gpu
上的数据指针,第三个是数据的大小,最后一个bool
类型表示是否存在cpu
数据。
关于SyncedHead head_
,前面是有定义的:
enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
这四个分别表明了数据的状态:未初始化、在CPU、在GPU、已同步。这个在具体实现的代码syncedmem.cpp
中可以看出来,比如这个函数to_cpu()
:
inline void SyncedMemory::to_cpu() {
switch (head_) {
case UNINITIALIZED:
CaffeMallocHost(&cpu_ptr_, size_);
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_);
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;
}
}
这个函数的目的是把数据传到cpu上,那如果数据状态是未初始化,就需要创建数据,并把head_
设为HEAD_AT_CPU
,同时把own_cpu_data_
设为true
;而如果数据在gpu上,则把数据从gpu复制到cpu的数据中,并把head_
设为SYNCED
,同时把own_cpu_data_
设为true
。
通过syncedmem
的控制,可以把训练过程中的数据在cpu和gpu之间同步,另外因为Blob
中使用的是shared_ptr
,二者结合,内存上也是一种优化的方式。
注:我所使用的caffe
版本较早,与最新版可能存在差异。
__EOF__
版权声明: 本博客所有文章除特别声明外,均采用 BY-NC-SA 许可协议。转载请注明出处!
本文链接: https://hackcv.com/posts/caffe-syncedmem-%E6%BA%90%E7%A0%81%E7%AE%80%E6%9E%90/