Caffe2 - C++ API
A deep learning, cross platform ML framework
pack_segments.cc
1 
17 #include "caffe2/operators/pack_segments.h"
18 
19 namespace caffe2 {
20 
21 template <>
22 template <typename T>
23 bool PackSegmentsOp<CPUContext>::DoRunWithType() {
24  return DispatchHelper<
25  TensorTypes2<char, int32_t, int64_t, float, std::string>,
26  T>::call(this, Input(DATA));
27 }
28 
29 template <>
30 template <typename T, typename Data_T>
31 bool PackSegmentsOp<CPUContext>::DoRunWithType2() {
32  const auto& data = Input(DATA);
33  const auto& lengths = Input(LENGTHS);
34  auto* output = Output(0);
35  Tensor<CPUContext>* presence_mask = nullptr;
36  if (return_presence_mask_) {
37  presence_mask = Output(1);
38  }
39 
40  CAFFE_ENFORCE(data.ndim() >= 1, "DATA should be at least 1-D");
41  CAFFE_ENFORCE(lengths.ndim() == 1, "LENGTH should be 1-D");
42 
43  // Find the length of the longest sequence.
44  const T* l = lengths.template data<T>();
45  T max_length = 0;
46  T total_length = 0;
47  for (T i = 0; i < lengths.dim(0); ++i) {
48  max_length = std::max(max_length, l[i]);
49  total_length += l[i];
50  }
51 
52  // Total lengths must be the same as data.dims(0)
53  CAFFE_ENFORCE_EQ(
54  data.dim(0),
55  total_length,
56  " PackSegments requires that the sum of the lengths ",
57  total_length,
58  " is equal to the first data dimension ",
59  data.dim(0));
60 
61  auto shape = data.dims(); // Shape of output is batch_size x max_len x ...
62  shape[0] = max_length;
63  shape.insert(shape.begin(), lengths.size());
64  output->Resize(shape);
65 
66  // create output tensor
67  auto* out = static_cast<char*>(output->raw_mutable_data(data.meta()));
68 
69  bool* presence_mask_data = nullptr;
70  if (return_presence_mask_) {
71  // Shape of presence is batch_size x max_len
72  std::vector<caffe2::TIndex> presence_shape{lengths.size(), max_length};
73  presence_mask->Resize(presence_shape);
74  presence_mask_data = presence_mask->template mutable_data<bool>();
75  }
76 
77  if (!data.dim(0)) {
78  // Return empty output (with the proper shape)
79  return true;
80  }
81 
82  // Do padding
83  if (output->template IsType<float>()) {
84  math::Set<float, CPUContext>(
85  output->size(),
86  padding_,
87  output->template mutable_data<float>(),
88  &context_);
89  }
90  if (return_presence_mask_) {
91  memset(presence_mask_data, (int)false, presence_mask->size());
92  }
93 
94  int block_size = data.size() / data.dim(0);
95  int block_bytesize = data.nbytes() / data.dim(0);
96  const auto* d = static_cast<const char*>(data.raw_data());
97  int start = 0;
98  for (int i = 0; i < lengths.dim(0); ++i) {
99  context_.template CopyItems<CPUContext, CPUContext>(
100  data.meta(),
101  l[i] * block_size,
102  d + block_bytesize * start,
103  out + block_bytesize * max_length * i);
104  if (return_presence_mask_) {
105  memset(presence_mask_data + max_length * i, (int)true, l[i]);
106  }
107  start += l[i];
108  }
109 
110  return true;
111 }
112 
113 template <>
114 template <typename T>
115 bool UnpackSegmentsOp<CPUContext>::DoRunWithType() {
116  return DispatchHelper<
117  TensorTypes2<char, int32_t, int64_t, float, std::string>,
118  T>::call(this, Input(DATA));
119 }
120 
121 template <>
122 template <typename T, typename Data_T>
123 bool UnpackSegmentsOp<CPUContext>::DoRunWithType2() {
124  const auto& data = Input(DATA);
125  const auto& lengths = Input(LENGTHS);
126  auto* output = Output(0);
127 
128  CAFFE_ENFORCE(data.ndim() >= 2, "DATA should be at least 2-D");
129  CAFFE_ENFORCE(lengths.ndim() == 1, "LENGTH should be 1-D");
130 
131  const T* l = lengths.template data<T>();
132 
133  T total_l = std::accumulate(l, l + lengths.dim(0), 0);
134 
135  auto shape = data.dims();
136  CAFFE_ENFORCE(
137  shape[0] == lengths.dim(0), "LENGTH should match DATA in dimension 0");
138  shape.erase(shape.begin());
139  shape[0] = total_l;
140  output->Resize(shape);
141  // create output tensor
142  auto* out = static_cast<char*>(output->raw_mutable_data(data.meta()));
143  if (!(data.dim(0) * data.dim(1))) {
144  return true;
145  }
146  int block_size = data.size() / (data.dim(0) * data.dim(1));
147  int block_bytesize = data.nbytes() / (data.dim(0) * data.dim(1));
148  const auto* d = static_cast<const char*>(data.raw_data());
149  int start = 0;
150  for (int i = 0; i < lengths.dim(0); ++i) {
151  context_.template CopyItems<CPUContext, CPUContext>(
152  data.meta(),
153  l[i] * block_size,
154  d + block_bytesize * data.dim(1) * i,
155  out + block_bytesize * start);
156  start += l[i];
157  }
158  return true;
159 }
160 
161 REGISTER_CPU_OPERATOR(PackSegments, PackSegmentsOp<CPUContext>);
162 REGISTER_CPU_OPERATOR(UnpackSegments, UnpackSegmentsOp<CPUContext>);
163 
164 OPERATOR_SCHEMA(PackSegments)
165  .NumInputs(2)
166  .NumOutputs(1, 2)
167  .SetDoc(
168  "Map N dim tensor to N+1 dim based on length blob. Sequences that \
169  are shorter than the longest sequence are padded with zeros.")
170  .Input(
171  0,
172  "lengths",
173  "1-d int/long tensor contains the length in each of the output.")
174  .Input(1, "tensor", "N dim Tensor.")
175  .Output(
176  0,
177  "packed_tensor",
178  "N + 1 dim Tensor"
179  "where dim(1) is the max length"
180  ", dim(0) is the batch size.")
181  .Output(
182  1,
183  "presence_mask",
184  "2 dim boolean tensor"
185  ", false where packed_tensor is padded, true otherwise.")
186  .Arg(
187  "pad_minf",
188  "Padding number in the packed segments. Use true to pad \
189  -infinity, otherwise pad zeros")
190  .Arg(
191  "return_presence_mask",
192  "bool whether to return presence mask, false by default");
193 OPERATOR_SCHEMA(UnpackSegments)
194  .NumInputs(2)
195  .NumOutputs(1)
196  .SetDoc("Map N+1 dim tensor to N dim based on length blob")
197  .Input(
198  0,
199  "lengths",
200  "1-d int/long tensor contains the length in each of the input.")
201  .Input(1, "tensor", "N+1 dim Tensor.")
202  .Output(0, "packed_tensor", "N dim Tensor");
203 
205  using GradientMakerBase::GradientMakerBase;
206  vector<OperatorDef> GetGradientDefs() override {
207  return SingleGradientDef(
208  "UnpackSegments",
209  "",
210  vector<string>{I(0), GO(0)},
211  vector<string>{GI(1)});
212  }
213 };
214 REGISTER_GRADIENT(PackSegments, GetPackSegmentsGradient);
215 
217  using GradientMakerBase::GradientMakerBase;
218  vector<OperatorDef> GetGradientDefs() override {
219  return SingleGradientDef(
220  "PackSegments", "", vector<string>{I(0), GO(0)}, vector<string>{GI(1)});
221  }
222 };
223 REGISTER_GRADIENT(UnpackSegments, GetUnpackSegmentsGradient);
224 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.
static vector< OperatorDef > SingleGradientDef(const Args &...args)
a helper function to allow one to create one single operator def, which is usually the case for many ...
Copyright (c) 2016-present, Facebook, Inc.