Note: This article is available in Chinese only. 本文暂无英文版本。
View original
背景#
这个 layer 和 reduce layer 有一些相似，就干脆一起看了。 作用是输入至少两个 blob，然后对每个 blob 中的元素做一些运算，最后得到一个 blob。
caffe 支持的运算有 “PROD”、“SUM”、“MAX” 三种。
顺便提一句，TensorRT 支持的要多一些：
1
2enum class ElementWiseOperation : int
3{
4 kSUM = 0, //!< Sum of the two elements.
5 kPROD = 1, //!< Product of the two elements.
6 kMAX = 2, //!< Maximum of the two elements.
7 kMIN = 3, //!< Minimum of the two elements.
8 kSUB = 4, //!< Substract the second element from the first.
9 kDIV = 5, //!< Divide the first element by the second.
10 kPOW = 6 //!< The first element to the power of the second element.
11};proto#
1
2message EltwiseParameter {
3 enum EltwiseOp {
4 PROD = 0;
5 SUM = 1;
6 MAX = 2;
7 }
8 optional EltwiseOp operation = 1 [default = SUM]; // element-wise operation
9 repeated float coeff = 2; // blob-wise coefficient for SUM operation
10
11 // Whether to use an asymptotically slower (for >2 inputs) but stabler method
12 // of computing the gradient for the PROD operation. (No effect for SUM op.)
13 optional bool stable_prod_grad = 3 [default = true];
14}proto 里面的 coeff 是针对 SUM 操作的，可以给每一个 bottom blob 一个加权系数；stable_prod_grad 是 backward 用的，不用管。
c++ 实现#
代码比较容易看懂，加了一些注释。有两个地方可以提一下：一个是 PROD 和 MAX 的做法，都是先求前两个，再把得到的结果和后面的 blob 进行运算（其实是很自然的操作…似乎也没什么可说的 orz）。
另外一个是 mask 这个变量，是在 MAX 操作时用来标记在哪个 bottom blob 取到了最大值，反向传播时要用。
详细代码
1
2template <typename Dtype>
3void EltwiseLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
4 const vector<Blob<Dtype>*>& top) {
5 for (int i = 1; i < bottom.size(); ++i) {
6 CHECK(bottom[i]->shape() == bottom[0]->shape());
7 }
8 // check所有的bottom blob的shape都一样. 至少存在两个bottom blob
9 top[0]->ReshapeLike(*bottom[0]);
10 // If max operation, we will initialize the vector index part.
11 if (this->layer_param_.eltwise_param().operation() ==
12 EltwiseParameter_EltwiseOp_MAX && top.size() == 1) {
13 max_idx_.Reshape(bottom[0]->shape());
14 }
15}
16
17template <typename Dtype>
18void EltwiseLayer<Dtype>::Forward_cpu(
19 const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
20 int* mask = NULL;
21 const Dtype* bottom_data_a = NULL;
22 const Dtype* bottom_data_b = NULL;
23 const int count = top[0]->count();
24 Dtype* top_data = top[0]->mutable_cpu_data();
25 switch (op_) {
26 case EltwiseParameter_EltwiseOp_PROD:
27 caffe_mul(count, bottom[0]->cpu_data(), bottom[1]->cpu_data(), top_data);
28 for (int i = 2; i < bottom.size(); ++i) {
29 caffe_mul(count, top_data, bottom[i]->cpu_data(), top_data);
30 }
31 // 先算前两个,然后把结果和后面的每一个blob(如果还有的话)做运算
32 break;
33 case EltwiseParameter_EltwiseOp_SUM:
34 caffe_set(count, Dtype(0), top_data);
35 // 初始化top data为0
36 // TODO(shelhamer) does BLAS optimize to sum for coeff = 1?
37 for (int i = 0; i < bottom.size(); ++i) {
38 caffe_axpy(count, coeffs_[i], bottom[i]->cpu_data(), top_data);
39 }
40 break;
41 // mask干啥用的???
42 // forward应该用不到,是backward求梯度需要知道在哪个位置得到了最大值
43 case EltwiseParameter_EltwiseOp_MAX:
44 // Initialize
45 mask = max_idx_.mutable_cpu_data();
46 caffe_set(count, -1, mask);
47 caffe_set(count, Dtype(-FLT_MAX), top_data);
48 // bottom 0 & 1
49 bottom_data_a = bottom[0]->cpu_data();
50 bottom_data_b = bottom[1]->cpu_data();
51 for (int idx = 0; idx < count; ++idx) {
52 if (bottom_data_a[idx] > bottom_data_b[idx]) {
53 top_data[idx] = bottom_data_a[idx]; // maxval
54 mask[idx] = 0; // maxid
55 } else {
56 top_data[idx] = bottom_data_b[idx]; // maxval
57 mask[idx] = 1; // maxid
58 }
59 }
60 // bottom 2++
61 for (int blob_idx = 2; blob_idx < bottom.size(); ++blob_idx) {
62 bottom_data_b = bottom[blob_idx]->cpu_data();
63 for (int idx = 0; idx < count; ++idx) {
64 if (bottom_data_b[idx] > top_data[idx]) {
65 top_data[idx] = bottom_data_b[idx]; // maxval
66 mask[idx] = blob_idx; // maxid
67 }
68 }
69 }
70 break;
71 default:
72 LOG(FATAL) << "Unknown elementwise operation.";
73 }
74}