caffe 源码学习笔记(7) slice layer

背景 

ocr组那边有个shuffle net 的网络,里面有个pytorch op叫chunk,转成的onnx对应的op是 split

作用是:

Split a tensor into a list of tensors, along the specified 'axis'. Lengths of the parts can be specified using argument 'split'. Otherwise, the tensor is split to equal sized parts.

然后发现这个op模型转换里不支持转到caffe的layer,于是想办法支持了一下. 发现是要转到caffe的slice layer.(caffe也有一个split layer,但是这个split layer是除了一个输出blob作为多个layer的输入时用的)

proto

 1
 2message SliceParameter {
 3  // The axis along which to slice -- may be negative to index from the end
 4  // (e.g., -1 for the last axis).
 5  // By default, SliceLayer concatenates blobs along the "channels" axis (1).
 6  optional int32 axis = 3 [default = 1];
 7  repeated uint32 slice_point = 2;
 8
 9  // DEPRECATED: alias for "axis" -- does not support negative indexing.
10  optional uint32 slice_dim = 1 [default = 1];
11}
12
13

看起来slice_dim和axis是新旧两种写法,slice_point应该就是切割点. 这个layer文档仍然是"to do"状态,因此只能看代码了.

c++实现

 1
 2template <typename Dtype>
 3void SliceLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
 4      const vector<Blob<Dtype>*>& top) {
 5  const SliceParameter& slice_param = this->layer_param_.slice_param();
 6  CHECK(!(slice_param.has_axis() && slice_param.has_slice_dim()))
 7      << "Either axis or slice_dim should be specified; not both.";
 8  slice_point_.clear();
 9  std::copy(slice_param.slice_point().begin(),
10      slice_param.slice_point().end(),
11      std::back_inserter(slice_point_));
12}
13

layerSetUp就是单纯把slice_point用成员函数来存储.

reshape 其实是比较重点的部分,注意不提供slice_point的默认情况. 其他部分很好懂,就是把切割点计算成每一部分切割线段的长度.

 1template <typename Dtype>
 2void SliceLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
 3      const vector<Blob<Dtype>*>& top) {
 4  const int num_axes = bottom[0]->num_axes();
 5  const SliceParameter& slice_param = this->layer_param_.slice_param();
 6  if (slice_param.has_slice_dim()) {
 7    slice_axis_ = static_cast<int>(slice_param.slice_dim());
 8    // Don't allow negative indexing for slice_dim, a uint32 -- almost
 9    // certainly unintended.
10    CHECK_GE(slice_axis_, 0) << "casting slice_dim from uint32 to int32 "
11        << "produced negative result; slice_dim must satisfy "
12        << "0 <= slice_dim < " << kMaxBlobAxes;
13    CHECK_LT(slice_axis_, num_axes) << "slice_dim out of range.";
14  } else {
15    slice_axis_ = bottom[0]->CanonicalAxisIndex(slice_param.axis());
16  }
17  vector<int> top_shape = bottom[0]->shape();
18  const int bottom_slice_axis = bottom[0]->shape(slice_axis_);
19  // bottom_slice_axis为切割的那个维度的总数
20  num_slices_ = bottom[0]->count(0, slice_axis_);
21  //  计算[0,slice_axis_)的count
22  slice_size_ = bottom[0]->count(slice_axis_ + 1);
23  //  计算[slice_axis_+1,num_axes()]的体积
24  int count = 0;
25  if (slice_point_.size() != 0) {
26    CHECK_EQ(slice_point_.size(), top.size() - 1);
27    //  n个点把一条线段切割成n+1(top.size())份
28    CHECK_LE(top.size(), bottom_slice_axis);
29    //  抽屉原理,保证每一个输出blob至少有一份.
30    int prev = 0;
31    vector<int> slices;
32    //  slices保存每一个切割的长度
33    for (int i = 0; i < slice_point_.size(); ++i) {
34      CHECK_GT(slice_point_[i], prev);
35      slices.push_back(slice_point_[i] - prev);
36      prev = slice_point_[i];
37    }
38    slices.push_back(bottom_slice_axis - prev);
39    for (int i = 0; i < top.size(); ++i) {
40      top_shape[slice_axis_] = slices[i];
41      top[i]->Reshape(top_shape);
42      count += top[i]->count();
43    }
44  } else {
45    // 如果不填写 slice_point,默认是把slice维度平均分给所有输出的blob
46    // 比如输入为[1,3,M,224],在M所在的维度所slice,输出为3个blob
47    //  那么就会得到三个shape为 [1,3,M/3,224]的blob,并且保证M%3==0; 
48    CHECK_EQ(bottom_slice_axis % top.size(), 0)
49        << "Number of top blobs (" << top.size() << ") should evenly "
50        << "divide input slice axis (" << bottom_slice_axis << ")";
51    top_shape[slice_axis_] = bottom_slice_axis / top.size();
52    for (int i = 0; i < top.size(); ++i) {
53      top[i]->Reshape(top_shape);
54      count += top[i]->count();
55    }
56  }
57  CHECK_EQ(count, bottom[0]->count());
58  if (top.size() == 1) {
59    top[0]->ShareData(*bottom[0]);
60    top[0]->ShareDiff(*bottom[0]);
61  }
62}
63

然后是forward,没什么好说的.

 1
 2template <typename Dtype>
 3void SliceLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
 4      const vector<Blob<Dtype>*>& top) {
 5  if (top.size() == 1) { return; }
 6  int offset_slice_axis = 0;
 7  const Dtype* bottom_data = bottom[0]->cpu_data();
 8  const int bottom_slice_axis = bottom[0]->shape(slice_axis_);
 9  for (int i = 0; i < top.size(); ++i) {
10    Dtype* top_data = top[i]->mutable_cpu_data();
11    const int top_slice_axis = top[i]->shape(slice_axis_);
12    for (int n = 0; n < num_slices_; ++n) {
13      const int top_offset = n * top_slice_axis * slice_size_;
14      const int bottom_offset =
15          (n * bottom_slice_axis + offset_slice_axis) * slice_size_;
16      caffe_copy(top_slice_axis * slice_size_,
17          bottom_data + bottom_offset, top_data + top_offset);
18    }
19    offset_slice_axis += top_slice_axis;
20  }
21}
22

Posts in this Series