跳过正文
  1. Posts/

caffe 源码学习笔记(10) eltwise layer

·988 字·2 分钟
99 - 这篇文章属于一个选集。

背景
#

这个 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}
99 - 这篇文章属于一个选集。

相关文章

caffe 源码学习笔记(6) reshape layer

·1851 字·4 分钟
背景 # 最近在魔改 TensorRT 的 caffe parser。 之前 caffe 模型转到 trt 模型时,有一个修改需要将 reshape layer 的 param 末尾补 1,比较繁琐,于是看了下 caffe 的 reshape layer 的实现。