caffe代码阅读6：SyncedMemory的j介绍与实现

一、SyncedMemory的作用简介

SyncedMemory类主要负责在GPU或者CPU上分配内存以及保持数据的同步作用。

SyncedMemory类主要应用在BLOB类中，我们可以在BLOB类中看出一些使用方法

比如：

[cpp]  view plain  copy

 

1. // blob中的reshape 的具体实现  
2. template <typename Dtype>  
3. void Blob<Dtype>::Reshape(const vector<int>& shape) {  
4.   CHECK_LE(shape.size(), kMaxBlobAxes); //是否小于规定的最大BLOB的维度(35维)  
5.   count_ = 1;  
6.   shape_.resize(shape.size());// 首先将大小设置为vector<int> shape_; 即新的形状数据的大小  
7.   if (!shape_data_ || shape_data_->size() < shape.size() * sizeof(int)) {  
8.     shape_data_.reset(new SyncedMemory(shape.size() * sizeof(int)));//  shared_ptr<SyncedMemory> shape_data_;  
9.   }  
10.   int* shape_data = static_cast<int*>(shape_data_->mutable_cpu_data());  
11.   for (int i = 0; i < shape.size(); ++i) {  
12.     // 检查形状数据是否合法  
13.     CHECK_GE(shape[i], 0);  
14.     CHECK_LE(shape[i], INT_MAX / count_) << "blob size exceeds INT_MAX";  
15.     // 计算数据个数  
16.     count_ *= shape[i];  
17.     // 复制shape到新的和旧的形状数据  
18.     shape_[i] = shape[i];  
19.     shape_data[i] = shape[i];  
20.   }  
21.   // 判断是否大于存储的容量  
22.   if (count_ > capacity_) {  
23.     capacity_ = count_;  
24.     // 重新分配内存  
25.     data_.reset(new SyncedMemory(capacity_ * sizeof(Dtype)));  
26.     diff_.reset(new SyncedMemory(capacity_ * sizeof(Dtype)));  
27.   }  
28. }  
29.   
30. // 注意这里的 gpu_data和mutable_gpu_data的区别是  
31. // mutable_gpu_data是设置了head_ 的而  
32. // gpu_data是不设置head_  
33. // 但是我并不明白为啥要有个mutable  
34. // 难道是因为设置了head_就可以实现互斥了?  
35. // 可能的解释是设置了head_可以保护gpu或者内存中的数据不被销毁  
36.   
37. template <typename Dtype>  
38. const int* Blob<Dtype>::gpu_shape() const {  
39.   CHECK(shape_data_);  
40.   // shared_ptr<SyncedMemory> shape_data_;  
41.   // 因此也分gpu_data和cpu_data  
42.   return (const int*)shape_data_->gpu_data();  
43. }  

另一个blob中使用的例子：

[cpp]  view plain  copy

 

1. // Update是计算data=-1 * diff + data  
2. template <typename Dtype>  
3. void Blob<Dtype>::Update() {  
4.   // We will perform update based on where the data is located.  
5.   switch (data_->head()) {  
6.   case SyncedMemory::HEAD_AT_CPU:  
7.     // perform computation on CPU  
8.     // axpby即alpha * x plus beta *y 这个含义,blas的函数命名真是见名知意  
9.     // template <> void caffe_axpy<float>(const int N, const float alpha, const float* X, float* Y) { cblas_saxpy(N, alpha, X, 1, Y, 1); }  
10.     // caffe_axpy计算的是Y=alpha * X + Y ，其中alpha=-1了这里  
11.     // 存储的时候用到了mutable_cpu_data，防止其他线程访问  
12.     caffe_axpy<Dtype>(count_, Dtype(-1),  
13.         static_cast<const Dtype*>(diff_->cpu_data()),  
14.         static_cast<Dtype*>(data_->mutable_cpu_data()));  
15.     break;  
16.   case SyncedMemory::HEAD_AT_GPU:  
17.   case SyncedMemory::SYNCED:  
18. #ifndef CPU_ONLY  
19.     // perform computation on GPU  
20.     // Y=alpha * X + Y ，其中alpha=-1了这里  
21.     caffe_gpu_axpy<Dtype>(count_, Dtype(-1),  
22.         static_cast<const Dtype*>(diff_->gpu_data()),  
23.         static_cast<Dtype*>(data_->mutable_gpu_data()));  
24. #else  
25.     NO_GPU;  
26. #endif  
27.     break;  
28.   default:  
29.     LOG(FATAL) << "Syncedmem not initialized.";  
30.   }  
31. }  

二、SyncedMemory类的详细介绍

1）构造函数

[cpp]  view plain  copy

 

1. // 构造函数  
2.   SyncedMemory()  
3.         // 一开始的时候就是未初始化  
4.       : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),  
5.         own_cpu_data_(false), own_gpu_data_(false), gpu_device_(-1) {}  
6.   explicit SyncedMemory(size_t size)  
7.       : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED),  
8.         own_cpu_data_(false), own_gpu_data_(false), gpu_device_(-1) {}  
9.    // 析构函数  
10.   ~SyncedMemory();  

2）成员变量

[cpp]  view plain  copy

 

1. // 指向内存的指针  
2.   void* cpu_ptr_;  
3.   void* gpu_ptr_;  
4.   // 数据大小  
5.   size_t size_;  
6.   // 同步状态  
7.   SyncedHead head_;  
8.   // 是否拥有cpu还是gpu数据  
9.   bool own_cpu_data_;  
10.   bool own_gpu_data_;  
11.   // 设备编号  
12.   int gpu_device_;  

3）成员函数

[cpp]  view plain  copy

 

1. // 获取CPUDATA  
2.   const void* cpu_data();  
3.   // 设置CPUDATA  
4.   void set_cpu_data(void* data);  
5.   // 获取GPUDATA  
6.   const void* gpu_data();  
7.   // 设置GPUDATA  
8.   void set_gpu_data(void* data);  
9.   // 获取互斥的CPU或者GPUDATA  
10.   void* mutable_cpu_data();  
11.   void* mutable_gpu_data();  
12.   // 枚举类型，未初始化，在CPU、在GPU、同步状态  
13.   enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };  
14.   // 获取数据的位置  
15.   SyncedHead head() { return head_; }  
16.   // 数据大小  
17.   size_t size() { return size_; }  
18. #ifndef CPU_ONLY  
19.   void async_gpu_push(const cudaStream_t& stream);  
20. #endif  
21.   
22.  private:  
23.  // 内部使用的到cpu还是gpu  
24.   void to_cpu();  
25.   void to_gpu();  

具体的实现如下：

三、SyncedMemory类的具体实现

[cpp]  view plain  copy

 

1. #include <cstring>  
2.   
3. #include "caffe/common.hpp"  
4. #include "caffe/syncedmem.hpp"  
5. #include "caffe/util/math_functions.hpp"  
6.   
7. namespace caffe {  
8. // 析构函数就是释放内存  
9. SyncedMemory::~SyncedMemory() {  
10.   if (cpu_ptr_ && own_cpu_data_) {  
11.     CaffeFreeHost(cpu_ptr_);  
12.   }  
13.   
14. #ifndef CPU_ONLY// 只要不是定义的CPU_ONLY的编译模式  
15.   if (gpu_ptr_ && own_gpu_data_) {  
16.     int initial_device;  
17.     // 获取可用设备  
18.     cudaGetDevice(&initial_device);  
19.     if (gpu_device_ != -1) {  
20.         // 当前所使用的设备  
21.       CUDA_CHECK(cudaSetDevice(gpu_device_));  
22.     }  
23.     // 释放当前  
24.     CUDA_CHECK(cudaFree(gpu_ptr_));  
25.     cudaSetDevice(initial_device);  
26.   }  
27. #endif  // CPU_ONLY  
28. }  
29.   
30. // 内部使用的  
31. // 如果当前未初始化，直接在内存分配空间  
32. // 如果在GPU上则复制到内存  
33. // 如果已经在内存则啥都不动  
34. inline void SyncedMemory::to_cpu() {  
35.   switch (head_) {  
36.   // 如果当前是未初始化，直接分配CPU上的内存  
37.   case UNINITIALIZED:  
38.     CaffeMallocHost(&cpu_ptr_, size_);  
39.     caffe_memset(size_, 0, cpu_ptr_);  
40.     head_ = HEAD_AT_CPU;  
41.     own_cpu_data_ = true;  
42.     break;  
43.   case HEAD_AT_GPU:  
44. #ifndef CPU_ONLY  
45.     // 如果当前数据在GPU，然后cpu_ptr为空  
46.     if (cpu_ptr_ == NULL) {  
47.         // 分配内存  
48.       CaffeMallocHost(&cpu_ptr_, size_);  
49.       own_cpu_data_ = true;  
50.     }  
51.     // 复制数据  
52.     caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_);  
53.     head_ = SYNCED;  
54. #else// CPU_ONLY模式当然只能报错了  
55.     NO_GPU;  
56. #endif  
57.     break;  
58.   case HEAD_AT_CPU:  
59.   case SYNCED:  
60.     break;  
61.   }  
62. }  
63.   
64. // 内部使用的  
65. // 如果当前未初始化直接在GPU分配内存  
66. // 如果当前在CPU，则在GPU上分配内存并且复制到GPU  
67. // 如果数据已经在GPU则啥也不做  
68. inline void SyncedMemory::to_gpu() {  
69. #ifndef CPU_ONLY  
70.   switch (head_) {  
71.   case UNINITIALIZED:  
72.     // 获取设备  
73.     CUDA_CHECK(cudaGetDevice(&gpu_device_));  
74.     // 在设备上分配内存  
75.     CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));  
76.     // 初始化为0  
77.     caffe_gpu_memset(size_, 0, gpu_ptr_);  
78.     head_ = HEAD_AT_GPU;  
79.     own_gpu_data_ = true;  
80.     break;  
81.   case HEAD_AT_CPU:  
82.     if (gpu_ptr_ == NULL) {  
83.       CUDA_CHECK(cudaGetDevice(&gpu_device_));  
84.       CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));  
85.       own_gpu_data_ = true;  
86.     }  
87.     caffe_gpu_memcpy(size_, cpu_ptr_, gpu_ptr_);  
88.     head_ = SYNCED;  
89.     break;  
90.   case HEAD_AT_GPU:  
91.   case SYNCED:  
92.     break;  
93.   }  
94. #else  
95.   NO_GPU;  
96. #endif  
97. }  
98.   
99. // 首先不管三七二十一将数据搞到内存上去  
100. // 然后获取cpu上的数据  
101. const void* SyncedMemory::cpu_data() {  
102.   to_cpu();  
103.   return (const void*)cpu_ptr_;  
104. }  
105.   
106.   
107. // 如果当前cpu_ptr_有内存上的数据则先释放  
108. // 然后再将地址给内部变量cpu_ptr_  
109. // 设置cpu上的数据  
110. void SyncedMemory::set_cpu_data(void* data) {  
111.   CHECK(data);  
112.   if (own_cpu_data_) {  
113.     CaffeFreeHost(cpu_ptr_);  
114.   }  
115.   cpu_ptr_ = data;  
116.   head_ = HEAD_AT_CPU;  
117.   own_cpu_data_ = false;  
118. }  
119.   
120. // 首先不管三七二十一将数据搞到GPU上去  
121. // 然后在获取gpu上的数据  
122. // 但是并没有改变head_的值(head_表明数据究竟在哪儿)  
123. const void* SyncedMemory::gpu_data() {  
124. #ifndef CPU_ONLY  
125.   to_gpu();  
126.   return (const void*)gpu_ptr_;  
127. #else  
128.   NO_GPU;  
129. #endif  
130. }  
131.   
132. // 如果当前gpu_ptr_有内存上的数据则先释放  
133. // 然后再将地址给内部变量gpu_ptr_  
134. // 设置gpu上的数据  
135. void SyncedMemory::set_gpu_data(void* data) {  
136. #ifndef CPU_ONLY  
137.   CHECK(data);  
138.   if (own_gpu_data_) {  
139.     int initial_device;  
140.     cudaGetDevice(&initial_device);  
141.     if (gpu_device_ != -1) {  
142.       CUDA_CHECK(cudaSetDevice(gpu_device_));  
143.     }  
144.     CUDA_CHECK(cudaFree(gpu_ptr_));  
145.     cudaSetDevice(initial_device);  
146.   }  
147.   gpu_ptr_ = data;  
148.   head_ = HEAD_AT_GPU;  
149.   own_gpu_data_ = false;  
150. #else  
151.   NO_GPU;  
152. #endif  
153. }  
154.   
155. // 首先不管三七二十一先数据搞到CPU上去  
156. // 然后返回互斥的cpu_ptr_指针  
157. // mutable_cpu_data与cpu_data的区别就是是否设置head  
158. void* SyncedMemory::mutable_cpu_data() {  
159.   to_cpu();  
160.   head_ = HEAD_AT_CPU;  
161.   return cpu_ptr_;  
162. }  
163.   
164.   
165. // 首先不管三七二十一先数据搞到GPU上去  
166. // 然后返回互斥的gpu_ptr_指针  
167. // mutable_gpu_data与gpu_data的区别就是是否设置head  
168. void* SyncedMemory::mutable_gpu_data() {  
169. #ifndef CPU_ONLY  
170.   to_gpu();  
171.   head_ = HEAD_AT_GPU;  
172.   return gpu_ptr_;  
173. #else  
174.   NO_GPU;  
175. #endif  
176. }  
177.   
178. #ifndef CPU_ONLY  
179. // 异步推送数据到gpu  
180. void SyncedMemory::async_gpu_push(const cudaStream_t& stream) {  
181.   CHECK(head_ == HEAD_AT_CPU);  
182.   if (gpu_ptr_ == NULL) {  
183.     CUDA_CHECK(cudaGetDevice(&gpu_device_));  
184.     CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));  
185.     own_gpu_data_ = true;  
186.   }  
187.   const cudaMemcpyKind put = cudaMemcpyHostToDevice;  
188.   CUDA_CHECK(cudaMemcpyAsync(gpu_ptr_, cpu_ptr_, size_, put, stream));  
189.   // Assume caller will synchronize on the stream before use  
190.   head_ = SYNCED;  
191. }  
192. #endif  
193.   
194. }  // namespace caffe  

四、总结

该类主要就是在内存分配空间以及在GPU上分配空间，并且负责同步数据，此外我看mutable_cpu_data和cpu_data 这两个函数的主要区别就是head_是否改变，至于这两个函数的命名上的mutable是有着互斥的含义的。究竟体现在哪儿，我的感受是，这里的mutable的体现主要是在调用了mutable_cpu_data之后强制设置了head_为HEAD_AT_CPU，从而保护了cpu上的数据。