Note: This article is available in Chinese only. 本文暂无英文版本。
View original
背景 #
OCR 组那边有个 shufflenet 的网络，里面有个 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}看起来 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}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}然后是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}