forked from BVLC/caffe
-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathreduction_layer.cpp
More file actions
129 lines (120 loc) · 4.21 KB
/
reduction_layer.cpp
File metadata and controls
129 lines (120 loc) · 4.21 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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
#include <vector>
#include "caffe/layers/reduction_layer.hpp"
#include "caffe/util/math_functions.hpp"
namespace caffe {
template <typename Dtype>
void ReductionLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top) {
op_ = this->layer_param_.reduction_param().operation();
}
template <typename Dtype>
void ReductionLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top) {
axis_ = bottom[0]->CanonicalAxisIndex(
this->layer_param_.reduction_param().axis());
// In the output, we'll keep all axes up to the reduction axis, but
// throw away any after that.
// Note: currently reducing along non-tail axes is not supported; otherwise,
// we'd need to also copy any axes following an "end_axis".
vector<int> top_shape(bottom[0]->shape().begin(),
bottom[0]->shape().begin() + axis_);
top[0]->Reshape(top_shape);
num_ = bottom[0]->count(0, axis_);
dim_ = bottom[0]->count(axis_);
CHECK_EQ(num_, top[0]->count());
if (op_ == ReductionParameter_ReductionOp_SUM ||
op_ == ReductionParameter_ReductionOp_MEAN) {
vector<int> sum_mult_shape(1, dim_);
sum_multiplier_.Reshape(sum_mult_shape);
caffe_set(dim_, Dtype(1), sum_multiplier_.mutable_cpu_data());
}
coeff_ = this->layer_param().reduction_param().coeff();
if (op_ == ReductionParameter_ReductionOp_MEAN) {
coeff_ /= dim_;
}
}
template <typename Dtype>
void ReductionLayer<Dtype>::Forward_cpu(
const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
const Dtype* bottom_data = bottom[0]->cpu_data();
const Dtype* mult_data = NULL;
if (sum_multiplier_.count() > 0) {
mult_data = sum_multiplier_.cpu_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:
*top_data = caffe_cpu_dot(dim_, mult_data, bottom_data);
break;
case ReductionParameter_ReductionOp_ASUM:
*top_data = caffe_cpu_asum(dim_, bottom_data);
break;
case ReductionParameter_ReductionOp_SUMSQ:
*top_data = caffe_cpu_dot(dim_, bottom_data, bottom_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_cpu_data();
caffe_scal(num_, coeff_, top_data);
}
}
template <typename Dtype>
void ReductionLayer<Dtype>::Backward_cpu(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]->cpu_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_cpu_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_set(dim_, bottom_coeff, bottom_diff);
break;
case ReductionParameter_ReductionOp_ASUM:
caffe_cpu_sign(dim_, bottom_data, bottom_diff);
caffe_scal(dim_, bottom_coeff, bottom_diff);
break;
case ReductionParameter_ReductionOp_SUMSQ:
caffe_cpu_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;
}
}
#ifdef CPU_ONLY
STUB_GPU(ReductionLayer);
#endif
INSTANTIATE_CLASS(ReductionLayer);
REGISTER_LAYER_CLASS(Reduction);
} // namespace caffe