Caffe2 - C++ API
A deep learning, cross platform ML framework
sequence_ops.h
1 
17 #ifndef CAFFE2_OPERATORS_SEQUENCE_OPS_H_
18 #define CAFFE2_OPERATORS_SEQUENCE_OPS_H_
19 
20 #include "caffe2/core/operator.h"
21 #include "caffe2/core/tensor.h"
22 #include "caffe2/utils/math.h"
23 
24 namespace caffe2 {
25 
26 template <class Context>
27 class GatherPaddingOp final : public Operator<Context> {
28  public:
29  USE_OPERATOR_CONTEXT_FUNCTIONS;
30  GatherPaddingOp(const OperatorDef& operator_def, Workspace* ws)
31  : Operator<Context>(operator_def, ws),
32  startPaddingWidth_(
33  OperatorBase::GetSingleArgument<int>("padding_width", 1)),
34  endPaddingWidth_(
35  OperatorBase::GetSingleArgument<int>("end_padding_width", -1)) {
36  CAFFE_ENFORCE_GE(startPaddingWidth_, 0);
37  if (endPaddingWidth_ < 0) {
38  endPaddingWidth_ = startPaddingWidth_;
39  }
40  }
41 
42  bool RunOnDevice() override {
43  if (startPaddingWidth_ == 0 && endPaddingWidth_ == 0) {
44  Output(0)->Resize(std::vector<TIndex>(0));
45  Output(0)->template mutable_data<TIndex>();
46  if (OutputSize() == 2) {
47  Output(1)->Resize(std::vector<TIndex>(0));
48  Output(1)->template mutable_data<TIndex>();
49  }
50  return true;
51  }
53  this, Input(0));
54  }
55 
56  template <typename T>
57  bool DoRunWithType() {
58  const auto& in = Input(0);
59  CAFFE_ENFORCE_GE(in.ndim(), 1);
60  const int32_t outer_size = in.dims()[0];
61  const auto block_size = in.size_from_dim(1);
62  const auto pad_width = startPaddingWidth_ + endPaddingWidth_;
63 
64  // if no lengths is provided, assume it is a single full-span entry
65  const int32_t* lengths_ptr = &outer_size;
66  int64_t lengths_size = 1;
67  if (InputSize() > 1) {
68  const auto& lengths = Input(1);
69  lengths_ptr = lengths.template data<int32_t>();
70  lengths_size = lengths.size();
71  }
72  std::vector<TIndex> padShape(in.dims().begin() + 1, in.dims().end());
73  // output will contain accumulator over paddings
74  Output(0)->Resize(padShape);
75  T* padding_start_ptr = Output(0)->template mutable_data<T>();
76  math::Set<T, Context>(block_size, 0.0, padding_start_ptr, &context_);
77 
78  // if no end_padding is provided, assume it's the same as start_padding
79  T* padding_end_ptr = padding_start_ptr;
80  if (OutputSize() == 2) {
81  Output(1)->Resize(padShape);
82  padding_end_ptr = Output(1)->template mutable_data<T>();
83  math::Set<T, Context>(block_size, 0.0, padding_end_ptr, &context_);
84  }
85  GatherPadding<T>(
86  outer_size,
87  lengths_size,
88  block_size,
89  pad_width,
90  in.template data<T>(),
91  lengths_ptr,
92  padding_start_ptr,
93  padding_end_ptr);
94  return true;
95  }
96 
97  private:
98  template <typename T>
99  void GatherPadding(
100  const int outer_size,
101  const int lengths_size,
102  const int block_size,
103  const int pad_width,
104  const T* in_ptr,
105  const int* lengths_ptr,
106  T* padding_start_ptr,
107  T* padding_end_ptr);
108 
109  int startPaddingWidth_;
110  int endPaddingWidth_;
111  // Scratch space required by the CUDA version
112  Tensor<Context> lengths_prefix_sum_buffer_;
113  Tensor<Context> lengths_prefix_sum_;
114 };
115 
116 template <class Context>
117 class RemovePaddingOp final : public Operator<Context> {
118  public:
119  USE_OPERATOR_CONTEXT_FUNCTIONS;
120  RemovePaddingOp(const OperatorDef& operator_def, Workspace* ws)
121  : Operator<Context>(operator_def, ws),
122  startPaddingWidth_(
123  OperatorBase::GetSingleArgument<int>("padding_width", 1)),
124  endPaddingWidth_(
125  OperatorBase::GetSingleArgument<int>("end_padding_width", -1)) {
126  CAFFE_ENFORCE_GE(startPaddingWidth_, 0);
127  if (endPaddingWidth_ < 0) {
128  endPaddingWidth_ = startPaddingWidth_;
129  }
130  }
131 
132  bool RunOnDevice() override {
133  if (startPaddingWidth_ == 0 && endPaddingWidth_ == 0) {
134  Output(0)->CopyFrom(Input(0), &context_);
135  if (OutputSize() == 2) {
136  Output(1)->CopyFrom(Input(1), &context_);
137  }
138  return true;
139  }
141  this, Input(0));
142  }
143 
144  template <typename T>
145  bool DoRunWithType();
146 
147  private:
148  int startPaddingWidth_;
149  int endPaddingWidth_;
150 
151  // Scratch space required by the CUDA version
152  Tensor<Context> lengths_prefix_sum_buffer_;
153  Tensor<Context> lengths_prefix_sum_;
154 };
155 
156 template <class Context>
157 class AddPaddingOp final : public Operator<Context> {
158  public:
159  USE_OPERATOR_CONTEXT_FUNCTIONS;
160  AddPaddingOp(const OperatorDef& operator_def, Workspace* ws)
161  : Operator<Context>(operator_def, ws),
162  startPaddingWidth_(
163  OperatorBase::GetSingleArgument<int>("padding_width", 1)),
164  endPaddingWidth_(
165  OperatorBase::GetSingleArgument<int>("end_padding_width", -1)) {
166  CAFFE_ENFORCE_GE(startPaddingWidth_, 0);
167  if (endPaddingWidth_ < 0) {
168  endPaddingWidth_ = startPaddingWidth_;
169  }
170  }
171 
172  bool RunOnDevice() override {
173  if (startPaddingWidth_ == 0 && endPaddingWidth_ == 0) {
174  Output(0)->CopyFrom(Input(0), &context_);
175  if (OutputSize() == 2) {
176  Output(1)->CopyFrom(Input(1), &context_);
177  }
178  return true;
179  }
181  this, Input(0));
182  }
183 
184  template <typename T>
185  bool DoRunWithType() {
186  const auto& in = Input(0);
187  CAFFE_ENFORCE_GE(in.ndim(), 1);
188  const int32_t outer_size = in.dims()[0];
189  const auto block_size = in.size_from_dim(1);
190 
191  // if no lengths is provided, assume it is a single full-span entry
192  const int32_t* lengths_ptr = nullptr;
193  int32_t lengths_size = 1;
194  if (InputSize() > 1) {
195  const auto& lengths = Input(1);
196  lengths_ptr = lengths.template data<int32_t>();
197  lengths_size = lengths.size();
198  }
199 
200  // fetch paddings
201  // input_size == 2 : pad with zeros
202  // input_size == 3 : start and end paddings are the same
203  // input_size == 4 : different start and end paddings
204  const T* padding_start_ptr = nullptr;
205  const T* padding_end_ptr = nullptr;
206  if (InputSize() >= 3) {
207  auto& padding_start = Input(2);
208  CAFFE_ENFORCE_EQ(block_size, padding_start.size());
209  padding_start_ptr = padding_start.template data<T>();
210  }
211  if (InputSize() == 4) {
212  auto& padding_end = Input(3);
213  CAFFE_ENFORCE_EQ(block_size, padding_end.size());
214  padding_end_ptr = padding_end.template data<T>();
215  } else {
216  padding_end_ptr = padding_start_ptr;
217  }
218 
219  auto* out = Output(0);
220  {
221  auto out_dims = in.dims();
222  out_dims[0] += (startPaddingWidth_ + endPaddingWidth_) * lengths_size;
223  out->Resize(std::move(out_dims));
224  }
225  const auto* in_ptr = in.template data<T>();
226  auto* out_ptr = out->template mutable_data<T>();
227 
228  return MakePadding<T>(
229  in_ptr,
230  out_ptr,
231  lengths_ptr,
232  lengths_size,
233  outer_size,
234  padding_start_ptr,
235  padding_end_ptr,
236  block_size);
237  }
238 
239  private:
240  template <typename T>
241  bool MakePadding(
242  const T* in_ptr,
243  T* out_ptr,
244  const int32_t* lengths_ptr,
245  int32_t lengths_size,
246  int32_t outer_size,
247  const T* padding_start_ptr,
248  const T* padding_end_ptr,
249  int64_t block_size);
250 
251  int startPaddingWidth_;
252  int endPaddingWidth_;
253 
254  // Scratch space required by the CUDA version
255  Tensor<Context> lengths_prefix_sum_buffer_;
256  Tensor<Context> lengths_prefix_sum_;
257 };
258 
259 template <class Context>
260 class PadEmptySamplesOp : public Operator<Context> {
261  public:
262  USE_OPERATOR_CONTEXT_FUNCTIONS;
263  PadEmptySamplesOp(const OperatorDef& operator_def, Workspace* ws)
264  : Operator<Context>(operator_def, ws) {}
265 
266  bool RunOnDevice() override;
267 };
268 
269 } // namespace caffe2
270 
271 #endif // CAFFE2_OPERATORS_SEQUENCE_OPS_H_
Tensor is the basic class in Caffe2 that stores a contiguous memory with its shape information...
Definition: tensor.h:109
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
Copyright (c) 2016-present, Facebook, Inc.