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