Caffe2 - C++ API
A deep learning, cross platform ML framework
sparse_to_dense_mask_op.h
1 
17 #ifndef CAFFE2_OPERATORS_SPARSE_TO_DENSE_MASK_OP_H_
18 #define CAFFE2_OPERATORS_SPARSE_TO_DENSE_MASK_OP_H_
19 
20 #include <algorithm>
21 #include <unordered_map>
22 #include <vector>
23 #include "caffe2/core/context.h"
24 #include "caffe2/core/operator.h"
25 #include "caffe2/core/tensor.h"
26 #include "caffe2/utils/math.h"
27 
28 namespace caffe2 {
29 
30 template <class Context>
31 class SparseToDenseMaskBase : public Operator<Context> {
32  public:
33  USE_OPERATOR_CONTEXT_FUNCTIONS;
34  SparseToDenseMaskBase(const OperatorDef& operator_def, Workspace* ws)
35  : Operator<Context>(operator_def, ws) {
36  std::vector<int64_t> mask =
37  OperatorBase::template GetRepeatedArgument<int64_t>("mask");
38  featuresCount_ = mask.size();
39 
40  CAFFE_ENFORCE(!mask.empty(), "mask can't be empty");
41  auto biggest = *std::max_element(mask.begin(), mask.end());
42  dense_.assign(std::min(kMaxDenseSize, biggest + 1), -1);
43  for (int i = 0; i < mask.size(); i++) {
44  int64_t id = mask[i];
45  CAFFE_ENFORCE_GE(id, 0, "Only positive IDs are allowed.");
46  if (id >= kMaxDenseSize) {
47  CAFFE_ENFORCE(sparse_.count(id) == 0, "Duplicated id: ", id);
48  sparse_[id] = i;
49  } else {
50  CAFFE_ENFORCE(dense_[id] == -1, "Duplicated id: ", id);
51  dense_[id] = i;
52  }
53  }
54  }
55 
56  protected:
57  const int64_t kMaxDenseSize = 1024 * 128;
58 
59  std::unordered_map<int64_t, int> sparse_;
60  std::vector<int> dense_;
61  int featuresCount_;
62 
63  inline int getFeatureIdx(int64_t id) const {
64  if (id >= kMaxDenseSize) {
65  const auto& iter = sparse_.find(id);
66  if (iter == sparse_.end()) {
67  return -1;
68  } else {
69  return iter->second;
70  }
71  } else {
72  return (id >= dense_.size()) ? -1 : dense_[id];
73  }
74  }
75 };
76 
77 template <class Context>
78 class SparseToDenseMaskOp : public SparseToDenseMaskBase<Context> {
79  public:
80  USE_OPERATOR_CONTEXT_FUNCTIONS;
81  SparseToDenseMaskOp(const OperatorDef& operator_def, Workspace* ws)
82  : SparseToDenseMaskBase<Context>(operator_def, ws) {
83  returnPresenceMask_ = OperatorBase::template GetSingleArgument<bool>(
84  "return_presence_mask", false);
85  maxSkippedSparseIndices_ =
86  OperatorBase::template GetSingleArgument<int32_t>(
87  "max_skipped_indices", kMaxSkippedSparseIndices);
88  }
89 
90  bool RunOnDevice() override {
92  this, Input(INDICES));
93  }
94 
95  template <typename TInd>
96  bool DoRunWithType() {
97  auto& sparse_indices = Input(INDICES);
98  CAFFE_ENFORCE_EQ(sparse_indices.ndim(), 1);
99  auto& sparse_values = Input(VALUES);
100  CAFFE_ENFORCE_GE(sparse_values.ndim(), 1);
101  CAFFE_ENFORCE_EQ(sparse_indices.size(), sparse_values.dim(0));
102  auto& default_value = Input(DEFAULT);
103  CAFFE_ENFORCE_EQ(default_value.ndim() + 1, sparse_values.ndim());
104  CAFFE_ENFORCE_EQ(default_value.size(), sparse_values.size_from_dim(1));
105  CAFFE_ENFORCE(sparse_values.meta() == default_value.meta());
106 
107  const TInd* sparse_indices_vec = sparse_indices.template data<TInd>();
108  const char* sparse_values_vec =
109  static_cast<const char*>(sparse_values.raw_data());
110  const void* default_val = default_value.raw_data();
111 
112  TIndex block_size = default_value.size();
113  size_t block_nbytes = default_value.nbytes();
114 
115  const int cols = this->featuresCount_;
116  int rows = -1;
117  int32_t sparse_indices_length = sparse_indices.dim32(0);
118  const int32_t* lengths_vec = nullptr;
119  auto* output = Output(OUTPUTVALUE);
120  Tensor<Context>* presence_mask = nullptr;
121  if (returnPresenceMask_) {
122  presence_mask = Output(PRESENCEMASK);
123  }
124  vector<TIndex> shape;
125  if (InputSize() == 4) {
126  auto& lengths = Input(LENGTHS);
127  CAFFE_ENFORCE_EQ(lengths.ndim(), 1);
128  lengths_vec = lengths.template data<int32_t>();
129  rows = lengths.dim32(0);
130  }
131  if (rows == -1) {
132  // if the LENGTHS is not set, the output will be a vector
133  rows = 1;
134  lengths_vec = &sparse_indices_length;
135  } else {
136  shape.push_back(rows);
137  }
138  shape.push_back(cols);
139  if (returnPresenceMask_) {
140  presence_mask->Resize(shape);
141  }
142  shape.insert(
143  shape.end(), default_value.dims().begin(), default_value.dims().end());
144  output->Resize(shape);
145 
146  // init
147  // TODO: consider unrolling CopyItems to make elemental types copy faster
148  char* output_data =
149  static_cast<char*>(output->raw_mutable_data(sparse_values.meta()));
150  for (int i = 0; i < cols * rows; i++) {
151  context_.template CopyItems<Context, Context>(
152  default_value.meta(),
153  block_size,
154  default_val,
155  output_data + i * block_nbytes);
156  }
157  bool* presence_mask_data = nullptr;
158  if (returnPresenceMask_) {
159  presence_mask_data = presence_mask->template mutable_data<bool>();
160  math::Set<bool, Context>(
161  rows * cols, false, presence_mask_data, &context_);
162  }
163 
164  int64_t offset = 0;
165  for (int r = 0; r < rows; r++) {
166  for (int c = 0; c < lengths_vec[r]; c++) {
167  const auto sparse_index = sparse_indices_vec[offset + c];
168  if (sparse_index < 0 ||
169  sparse_index >= std::numeric_limits<TInd>::max()) {
170  CAFFE_ENFORCE_LT(
171  ++skippedSparseIndices_,
172  maxSkippedSparseIndices_,
173  "Too many sparse indices skipped");
174  continue;
175  }
176  int idx = this->getFeatureIdx(sparse_index);
177  if (idx != -1) {
178  context_.template CopyItems<Context, Context>(
179  sparse_values.meta(),
180  block_size,
181  sparse_values_vec + (offset + c) * block_nbytes,
182  output_data + (r * cols + idx) * block_nbytes);
183  if (returnPresenceMask_) {
184  presence_mask_data[r * cols + idx] = true;
185  }
186  }
187  }
188  offset += lengths_vec[r];
189  }
190 
191  return true;
192  }
193 
194  private:
195  static const uint32_t kMaxSkippedSparseIndices = 5;
196 
197  bool returnPresenceMask_;
198  uint32_t maxSkippedSparseIndices_ = 0;
199  uint32_t skippedSparseIndices_ = 0;
200 
201  INPUT_TAGS(INDICES, VALUES, DEFAULT, LENGTHS);
202  OUTPUT_TAGS(OUTPUTVALUE, PRESENCEMASK);
203 };
204 
205 template <class Context>
207  public:
208  USE_OPERATOR_CONTEXT_FUNCTIONS;
209  SparseToDenseMaskGradientOp(const OperatorDef& operator_def, Workspace* ws)
210  : SparseToDenseMaskBase<Context>(operator_def, ws) {}
211 
212  bool RunOnDevice() override {
214  this, Input(INDICES));
215  }
216 
217  template <typename TInd>
218  bool DoRunWithType() {
219  auto& sparse_indices = Input(INDICES);
220  CAFFE_ENFORCE_EQ(sparse_indices.ndim(), 1);
221  auto& gradient_output = Input(GOUTPUT);
222 
223  TIndex block_size = gradient_output.size_from_dim(1);
224  size_t block_nbytes = gradient_output.itemsize() * block_size;
225 
226  const int cols = this->featuresCount_;
227  int rows = -1;
228  int iter_offset = 1;
229  int32_t default_length = sparse_indices.dim32(0);
230  const int32_t* lengths_vec = nullptr;
231  auto* output = Output(GVALUES);
232  vector<TIndex> shape;
233  if (InputSize() > LENGTHS) {
234  // if the LENGTHS is set, the gradient_output has dim:
235  // lengths * mask.size() * feature_dim
236  auto& lengths = Input(LENGTHS);
237  lengths_vec = lengths.template data<int32_t>();
238  rows = lengths.dim32(0);
239  CAFFE_ENFORCE_EQ(lengths.ndim(), 1);
240  CAFFE_ENFORCE_GE(gradient_output.ndim(), 2);
241  CAFFE_ENFORCE_EQ(gradient_output.dim(0), rows);
242  CAFFE_ENFORCE_EQ(gradient_output.dim(1), cols);
243  block_nbytes /= gradient_output.dim(1);
244  block_size /= gradient_output.dim(1);
245  iter_offset += 1;
246  }
247  if (rows == -1) {
248  // if the LENGTHS is not set, the gradient_output has dim:
249  // mask.size() * feature_dim
250  rows = 1;
251  lengths_vec = &default_length;
252  CAFFE_ENFORCE_GE(gradient_output.ndim(), 1);
253  CAFFE_ENFORCE_EQ(gradient_output.dim(0), cols);
254  }
255  shape.push_back(default_length);
256  // insert feature_dim
257  shape.insert(
258  shape.end(),
259  gradient_output.dims().begin() + iter_offset,
260  gradient_output.dims().end());
261  output->Resize(shape);
262 
263  const TInd* sparse_indices_vec = sparse_indices.template data<TInd>();
264  const char* gradient_output_vec =
265  static_cast<const char*>(gradient_output.raw_data());
266 
267  char* output_data =
268  static_cast<char*>(output->raw_mutable_data(gradient_output.meta()));
269  math::Set<char, Context>(
270  default_length * gradient_output.itemsize(), 0, output_data, &context_);
271 
272  int32_t offset = 0;
273  // SparseToDenseMask is not injective; gradient_used records
274  // if the gradient is used for other input value from the same row
275  vector<bool> gradient_used(cols, false);
276  for (int r = 0; r < rows; r++) {
277  std::fill(gradient_used.begin(), gradient_used.end(), false);
278  for (int c = lengths_vec[r] - 1; c >= 0; c--) {
279  int idx = this->getFeatureIdx(sparse_indices_vec[offset + c]);
280  if (idx != -1 && !gradient_used[idx]) {
281  gradient_used[idx] = true;
282  context_.template CopyItems<Context, Context>(
283  gradient_output.meta(),
284  block_size,
285  gradient_output_vec + (r * cols + idx) * block_nbytes,
286  output_data + (offset + c) * block_nbytes);
287  }
288  }
289  offset += lengths_vec[r];
290  }
291  return true;
292  }
293 
294  private:
295  INPUT_TAGS(INDICES, GOUTPUT, LENGTHS);
296  OUTPUT_TAGS(GVALUES);
297 };
298 
299 } // namespace caffe2
300 
301 #endif // CAFFE2_OPERATORS_SPARSE_TO_DENSE_MASK_OP_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
void Resize(Ts...dim_source)
Resizes a tensor.
Definition: tensor.h:304
Copyright (c) 2016-present, Facebook, Inc.