-
Notifications
You must be signed in to change notification settings - Fork 8
Expand file tree
/
Copy pathreduction_layer.cu
More file actions
91 lines (85 loc) · 2.85 KB
/
reduction_layer.cu
File metadata and controls
91 lines (85 loc) · 2.85 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
#include <vector>
#include "caffe/layers/reduction_layer.hpp"
#include "caffe/util/math_functions.hpp"
namespace caffe {
template <typename Dtype>
void ReductionLayer<Dtype>::Forward_gpu(
const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
const Dtype* bottom_data = bottom[0]->gpu_data();
const Dtype* mult_data = NULL;
if (sum_multiplier_.count() > 0) {
mult_data = sum_multiplier_.gpu_data();
}
Dtype* top_data = top[0]->mutable_cpu_data();
for (int i = 0; i < num_; ++i) {
switch (op_) {
case ReductionParameter_ReductionOp_SUM:
case ReductionParameter_ReductionOp_MEAN:
caffe_gpu_dot(dim_, mult_data, bottom_data, top_data);
break;
case ReductionParameter_ReductionOp_ASUM:
caffe_gpu_asum(dim_, bottom_data, top_data);
break;
case ReductionParameter_ReductionOp_SUMSQ:
caffe_gpu_dot(dim_, bottom_data, bottom_data, top_data);
break;
default:
LOG(FATAL) << "Unknown reduction op: "
<< ReductionParameter_ReductionOp_Name(op_);
}
bottom_data += dim_;
++top_data;
}
if (coeff_ != Dtype(1)) {
// Reset the top_data pointer.
top_data = top[0]->mutable_gpu_data();
caffe_gpu_scal(num_, coeff_, top_data);
}
}
template <typename Dtype>
void ReductionLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype>*>& top,
const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
if (!propagate_down[0]) { return; }
// Get bottom_data, if needed.
const Dtype* bottom_data = NULL;
switch (op_) {
// Operations that don't need bottom_data
case ReductionParameter_ReductionOp_SUM:
case ReductionParameter_ReductionOp_MEAN:
break;
// Operations that need bottom_data
case ReductionParameter_ReductionOp_ASUM:
case ReductionParameter_ReductionOp_SUMSQ:
bottom_data = bottom[0]->gpu_data();
break;
default:
LOG(FATAL) << "Unknown reduction op: "
<< ReductionParameter_ReductionOp_Name(op_);
}
const Dtype* top_diff = top[0]->cpu_diff();
Dtype* bottom_diff = bottom[0]->mutable_gpu_diff();
for (int i = 0; i < num_; ++i) {
const Dtype bottom_coeff = (*top_diff) * coeff_;
switch (op_) {
case ReductionParameter_ReductionOp_SUM:
case ReductionParameter_ReductionOp_MEAN:
caffe_gpu_set(dim_, bottom_coeff, bottom_diff);
break;
case ReductionParameter_ReductionOp_ASUM:
caffe_gpu_sign(dim_, bottom_data, bottom_diff);
caffe_gpu_scal(dim_, bottom_coeff, bottom_diff);
break;
case ReductionParameter_ReductionOp_SUMSQ:
caffe_gpu_scale(dim_, 2 * bottom_coeff, bottom_data, bottom_diff);
break;
default:
LOG(FATAL) << "Unknown reduction op: "
<< ReductionParameter_ReductionOp_Name(op_);
}
bottom_data += dim_;
bottom_diff += dim_;
++top_diff;
}
}
INSTANTIATE_LAYER_GPU_FUNCS(ReductionLayer);
} // namespace caffe