Caffe 源码 —— blob.hpp/cpp
文章目录
syncedmem.hpp/cpp
在介绍 Blob 之前得先说一说 syncedmem.hpp/cpp,syncedmem 文件中
定义了用于数据 CPU 和 GPU 之间的数据同步的 SyncedMemory 类,而这也是后面网络参数,梯度数据传输的基础。
主要包含了:
- CaffeMallocHost 内存申请函数
- CaffeFreeHost 内存释放函数
- SyncedMemory 类
1. CaffeMallocHost 函数
inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
#ifndef CPU_ONLY // 如果是 GPU 模式就调用 cudaMallocHost
if (Caffe::mode() == Caffe::GPU) {
CUDA_CHECK(cudaMallocHost(ptr, size));
*use_cuda = true;
return;
}
#endif
#ifdef USE_MKL // 使用 mkl 库则使用 mkl_malloc
*ptr = mkl_malloc(size ? size:1, 64);
#else
*ptr = malloc(size); // 正常 CPU 就直接使用 malloc
#endif
*use_cuda = false;
CHECK(*ptr) << "host allocation of size " << size << " failed";
}
可以看到上述函数根据当前环境的不同使用了不同的内存申请函数来完成为数据申请内存,CPU 模式下就直接用了 malloc 函数。
2. CaffeFreeHost 函数
与第 1 部分对应,我们申请了内存就要负责将资源释放掉。
inline void CaffeFreeHost(void* ptr, bool use_cuda) {
#ifndef CPU_ONLY // GPU 模式
if (use_cuda) {
CUDA_CHECK(cudaFreeHost(ptr));
return;
}
#endif
#ifdef USE_MKL // 使用 mkl 库
mkl_free(ptr);
#else
free(ptr); // CPU 模式
#endif
}
3. SyncedMemory 类
主要成员变量 | 主要成员函数 |
---|---|
void* cpu_ptr_; | void to_cpu(); |
void* gpu_ptr_; | void to_gpu(); |
bool own_cpu_data_; | void set_cpu_data(void* data); |
bool own_gpu_data_; | void set_gpu_data(void* data); |
SyncedHead head_; | —— |
这个类主要的作用就是同步 CPU 与 GPU 之间的数据,主要是通过上面几个主要的成员变量来进行控制,SyncedHead head_
是一个枚举类型。
enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
状态机变量,表示4种状态:未初始化、CPU数据有效、GPU数据有效、已同步。 以 to_cpu()
成员函数为例。
inline void SyncedMemory::to_cpu() {
check_device();
switch (head_) {
case UNINITIALIZED: // 如果未分配过内存(构造函数后就是这个状态)
CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_); // to_CPU时为CPU分配内存
caffe_memset(size_, 0, cpu_ptr_); // 数据清零
head_ = HEAD_AT_CPU; // 指示CPU更新了数据
own_cpu_data_ = true;
break;
case HEAD_AT_GPU: // 如果GPU侧更新过数据,则同步到CPU
#ifndef CPU_ONLY
if (cpu_ptr_ == NULL) { // 如果CPU侧没分配过内存,分配内存
CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
own_cpu_data_ = true;
}
caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_); // 数据同步
head_ = SYNCED; // 指示CPU和GPU数据已同步一致
#else
NO_GPU;
#endif
break;
case HEAD_AT_CPU: // 如果CPU数据是最新的,不操作
case SYNCED: // 如果CPU和GPU数据都是最新的,不操作
break;
}
}
可以看到如果当前数据的状态是未初始化,则通过调用内存分配函数对数据进行内存申请,并置成 0,最后赋到 cpu_ptr_
返回,这里如果是经由 to_cpu()
进行初始化的数据,有一个标志位 own_cpu_data_
会被设成 true
,表示当前拥有数据的所有权,这个所有权表示我们申请它同样要负责释放它。以当前数据在 CPU 上为例,如果当前数据需要同步到 GPU 上,则需要使用 caffe_gpu_memcpy()
来完成数据同步工作并表示数据同步已完成。
void SyncedMemory::set_cpu_data(void* data) {
check_device();
CHECK(data);
if (own_cpu_data_) { // 如果自己分配过内存,先释放,换外部指定数据
CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
}
cpu_ptr_ = data; // 直接指向外部数据
head_ = HEAD_AT_CPU; // 指示CPU侧更新了数据
own_cpu_data_ = false; // 指示数据来源于外部
}
主要体现在 set_cpu_data()
这个成员函数中,当我们从外部输入data
,这里的数据可以是网络参数也可以是参数梯度,如果当前我们拥有数据的所有权就需要将当前申请的数据释放掉,然后再将 cpu_ptr_
指向输入的数据 data
,完成数据设置,这时标志位 own_cpu_data_
会被设成 false
。
其他的成员函数还有 mutable_cpu_data()
这个也是直接调用了 to_cpu()
函数,在后面的 blob.cpp 中会有调用。
blob.hpp/cpp
谈到 Blob,它是 caffe 框架中数据的传输格式,是一个四维矩阵,分别对应:N,C,H,W,即 批大小,通道数, 特征图的高, 特征图的宽,并且提供了数据 data 和对应梯度 diff 的操作方式,为后面网络的训练提供基础。
本文件主要是定义了 Blob 的类,其与 SyncedMemory
类构成一种委托的关系,即 blob 类里包含了 SyncedMemory 类的指针。
shared_ptr<SyncedMemory> data_; // 存放指向 data 的指针
shared_ptr<SyncedMemory> diff_; // 存放指向 diff 的指针
shared_ptr<SyncedMemory> shape_data_; // 存放维度信息的指针
vector<int> shape_; // 形状信息
int count_; // 存放有效元素数目
int capacity_; // 存放 Blob 容器的容量信息
在 blob 类构造函数中主要就是在内部先初始化容量信息 capacity_
,然后调用了 void Reshape();
函数,其功能是改变 blob 的维度,如果有需要需要重新申请更大的内存。根据输入格式的不同, void Reshape()
函数重载了3种,分别是:
void Reshape(const int num, const int channels, const int height,const int width);
void Reshape(const BlobShape& shape);
void ReshapeLike(const Blob& other);
不管哪种形式最后都是统一转化并调用:
void Reshape(const vector<int>& shape);
来完成上述功能,下面是 Reshape 的函数定义。
template <typename Dtype>
void Blob<Dtype>::Reshape(const vector<int>& shape) {
CHECK_LE(shape.size(), kMaxBlobAxes); // 先检查输入数据维数有没有大于最大维数,合法是4
count_ = 1; // 用于计算元素总数 count_=num*channels*height*width, 初始输入是 1x3x48x48
shape_.resize(shape.size()); // 成员变量维度也被重置,这里是 vector 的 resize,shape_初始是 0 ,运行后变成 [0,0,0,0]
if (!shape_data_ || shape_data_->size() < shape.size() * sizeof(int)) { // sizeof(int) 为4, shape.size() 为4
shape_data_.reset(new SyncedMemory(shape.size() * sizeof(int))); // shape_data_ 的 size_ 为 16
}
int* shape_data = static_cast<int*>(shape_data_->mutable_cpu_data()); // 指针格式转化
for (int i = 0; i < shape.size(); ++i) {
CHECK_GE(shape[i], 0); // 保证每一个维度尺寸都 >=0
if (count_ != 0) {
// 保证 count_ 不溢出,应为下一次 count_ 的数量就是 shape[i] * count_
CHECK_LE(shape[i], INT_MAX / count_) << "blob size exceeds INT_MAX";
}
count_ *= shape[i]; // count_ 累乘
shape_[i] = shape[i]; // 为成员变量赋值 shape_[i] 分别为 1 3 48 48
shape_data[i] = shape[i];
}
if (count_ > capacity_) { //如果当前成员大于分配的空间容量,初始的 capacity_ 为 0
capacity_ = count_; // 扩容,重新分配 data_ 和 diff_ 的空间
data_.reset(new SyncedMemory(capacity_ * sizeof(Dtype))); // data_ 要分配 6912 * sizeof(Dtype) 个字节的内存大小
diff_.reset(new SyncedMemory(capacity_ * sizeof(Dtype))); // diff_ 要分配 6912 * sizeof(Dtype) 个字节的内存大小
}
}
可以看到是根据输入数据得 shape 信息与当前分配的空间容量做比较,然后重新分配对应的空间来存储数据。基于上面的构造函数,然后在 Blob 类中,使用 Update() 来更新网络参数,即 data=data-diff
来更新。
// Update() 函数用于网络参数 Blob 的更新
template <typename Dtype>
void Blob<Dtype>::Update() {
// We will perform update based on where the data is located.
switch (data_->head()) { // data 在哪就在哪更新
case SyncedMemory::HEAD_AT_CPU: // data 位于 CPU 端
// perform computation on CPU
// 执行在 CPU 上的计算 data_[i] = data_[i]-diff[i]
caffe_axpy<Dtype>(count_, Dtype(-1),
static_cast<const Dtype*>(diff_->cpu_data()),
static_cast<Dtype*>(data_->mutable_cpu_data()));
break;
case SyncedMemory::HEAD_AT_GPU: // data 位于 GPU 端,或者 CPU/GPU已同步
case SyncedMemory::SYNCED:
#ifndef CPU_ONLY
// perform computation on GPU
caffe_gpu_axpy<Dtype>(count_, Dtype(-1),
static_cast<const Dtype*>(diff_->gpu_data()),
static_cast<Dtype*>(data_->mutable_gpu_data()));
#else
NO_GPU;
#endif
break;
default:
LOG(FATAL) << "Syncedmem not initialized.";
}
}
这里的 caffe_axpy()
则是调用了标准 cblas 加速库的东西,同理,还有其他计算 data L1,L2范数的函数定义,以及它们乘以一个标量的函数,均是通过 data_->head()
的情况来调用加速库进行相关计算。
余下的就是一些对 blob 数据结构的查询功能,例如某一维的尺寸,对输入 index 的转化来查询到想对应的数据,还有就是数据的 Offset, 根据头指针及维度信息来查询到维度对应的实际存储信息的位置等。
其他的数据操作,访问 cpu 数据啊,读写 cpu 数据啊,都是直接调用的 SyncedMemory 类。就这样构建起了整个 caffe 框架的数据存储结构!