Note: This article is available in Chinese only. 本文暂无英文版本。
View original
背景#
其实没什么背景,继续啃 caffe 代码而已 2333。
reduce layer 其实就是做 reduce 操作,把一个任意 shape 的 blob 通过某种运算变成一个 scalar。
caffe 目前支持求和(SUM)、绝对值的和(ASUM)、平方和(SUMSQ),以及对得到的 scalar 的总数求平均的求和(MEAN)。
说句题外话,TensorRT 支持的操作是求和、求积、max、min 和 ave,还是有一些 gap 的。
proto#
先看 proto。
1
2message ReductionParameter {
3 enum ReductionOp {
4 SUM = 1;
5 ASUM = 2;
6 SUMSQ = 3;
7 MEAN = 4;
8 }
9
10 optional ReductionOp operation = 1 [default = SUM]; // reduction operation
11
12 // The first axis to reduce to a scalar -- may be negative to index from the
13 // end (e.g., -1 for the last axis).
14 // (Currently, only reduction along ALL "tail" axes is supported; reduction
15 // of axis M through N, where N < num_axes - 1, is unsupported.)
16 // Suppose we have an n-axis bottom Blob with shape:
17 // (d0, d1, d2, ..., d(m-1), dm, d(m+1), ..., d(n-1)).
18 // If axis == m, the output Blob will have shape
19 // (d0, d1, d2, ..., d(m-1)),
20 // and the ReductionOp operation is performed (d0 * d1 * d2 * ... * d(m-1))
21 // times, each including (dm * d(m+1) * ... * d(n-1)) individual data.
22 // If axis == 0 (the default), the output Blob always has the empty shape
23 // (count 1), performing reduction across the entire input --
24 // often useful for creating new loss functions.
25 optional int32 axis = 2 [default = 0];
26
27 optional float coeff = 3 [default = 1.0]; // coefficient for output
28}operation 不用说,表示要进行哪种 reduce 操作。 coeff 也比较简单,就是最后结果中每个元素都可以乘一个额外的系数,默认是 1。
axis 的含义是,从 axis 这个维度开始做 reduce。
比如 blob 的 shape 是 (a,b,c,axis,d,e),那么就需要对 (axis,d,e) 这部分做 a*b*c 次 reduce,得到 a*b*c 个标量。
需要注意,目前只支持从某个维度一直到最后的维度都去做 reduce,不支持中间的几个维度去做 reduce。
c++ 代码#
头文件中有两个成员变量 num_、dim_ 值得一说。
1
2template <typename Dtype>
3class ReductionLayer : public Layer<Dtype> {
4 public:
5// ...省略无关部分
6
7 /// @brief the reduction operation performed by the layer
8 ReductionParameter_ReductionOp op_;
9 /// @brief a scalar coefficient applied to all outputs
10 Dtype coeff_;
11 /// @brief the index of the first input axis to reduce
12 int axis_;
13 /// @brief the number of reductions performed
14 int num_;
15 /// @brief the input size of each reduction
16 int dim_;
17 /// @brief a helper Blob used for summation (op_ == SUM)
18 Blob<Dtype> sum_multiplier_;
19};num_ 表示的是要做 reduce 操作的次数,dim_ 表示的是每次 reduce 的维度部分的 count。
Reshape 部分主要就是计算了 num_、dim_、coeff_ 这几个成员变量。
1
2template <typename Dtype>
3void ReductionLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
4 const vector<Blob<Dtype>*>& top) {
5 axis_ = bottom[0]->CanonicalAxisIndex(
6 this->layer_param_.reduction_param().axis());
7 // In the output, we'll keep all axes up to the reduction axis, but
8 // throw away any after that.
9 // Note: currently reducing along non-tail axes is not supported; otherwise,
10 // we'd need to also copy any axes following an "end_axis".
11 vector<int> top_shape(bottom[0]->shape().begin(),
12 bottom[0]->shape().begin() + axis_);
13 top[0]->Reshape(top_shape);
14 // count都是左闭右开,一个参数就是到末尾
15 num_ = bottom[0]->count(0, axis_);
16 // num_表示的是要做reduce操作的次数
17 // range [0,axis_)
18 dim_ = bottom[0]->count(axis_);
19 // dim_表示的是reduce操作部分的count
20 // range [axis_,num_axes())
21 CHECK_EQ(num_, top[0]->count());
22 if (op_ == ReductionParameter_ReductionOp_SUM ||
23 op_ == ReductionParameter_ReductionOp_MEAN) {
24 vector<int> sum_mult_shape(1, dim_);
25 // sum_mult_shape[0] = dim_
26 sum_multiplier_.Reshape(sum_mult_shape);
27 caffe_set(dim_, Dtype(1), sum_multiplier_.mutable_cpu_data());
28 // sum_multiplier_是dim_长度的vector,值都为1
29 }
30 coeff_ = this->layer_param().reduction_param().coeff();
31 if (op_ == ReductionParameter_ReductionOp_MEAN) {
32 coeff_ /= dim_;
33 }
34}注意到求和的计算是构造了一个 dim_ 长度的、值全为 1 的向量,和 bottom 做点积实现的。
1
2template <typename Dtype>
3void ReductionLayer<Dtype>::Forward_cpu(
4 const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
5 const Dtype* bottom_data = bottom[0]->cpu_data();
6 const Dtype* mult_data = NULL;
7 if (sum_multiplier_.count() > 0) {
8 mult_data = sum_multiplier_.cpu_data();
9 }
10 Dtype* top_data = top[0]->mutable_cpu_data();
11 // 得到num_个结果
12 for (int i = 0; i < num_; ++i) {
13 switch (op_) {
14 case ReductionParameter_ReductionOp_SUM:
15 case ReductionParameter_ReductionOp_MEAN:
16 *top_data = caffe_cpu_dot(dim_, mult_data, bottom_data);
17 break;
18 case ReductionParameter_ReductionOp_ASUM:
19 *top_data = caffe_cpu_asum(dim_, bottom_data);
20 break;
21 case ReductionParameter_ReductionOp_SUMSQ:
22 *top_data = caffe_cpu_dot(dim_, bottom_data, bottom_data);
23 break;
24 default:
25 LOG(FATAL) << "Unknown reduction op: "
26 << ReductionParameter_ReductionOp_Name(op_);
27 }
28 bottom_data += dim_;
29 ++top_data;
30 }
31 // coeff_是每一个位置的系数,默认为1
32 // 此时才真的把系数放在每一个元素上
33 if (coeff_ != Dtype(1)) {
34 // Reset the top_data pointer.
35 top_data = top[0]->mutable_cpu_data();
36 caffe_scal(num_, coeff_, top_data);
37 }
38}