Caffe2 - C++ API
A deep learning, cross platform ML framework
unique_ops.cc
1 
17 #include "caffe2/operators/unique_ops.h"
18 
19 #include <cmath>
20 
21 namespace caffe2 {
22 
23 template <>
24 template <typename T>
25 bool UniqueOp<CPUContext>::DoRunWithType() {
26  auto& inputTensor = Input(0);
27  // use dim32 to enforce that it's fine to have remapping of type int
28  int N = inputTensor.dim32(0);
29  CAFFE_ENFORCE_EQ(inputTensor.dim(), 1, "Input should be a vector");
30 
31  int* remapping = nullptr;
32  if (REMAPPING < OutputSize()) {
33  auto* remappingTensor =
34  Output(REMAPPING, inputTensor.sizes(), at::dtype<int>());
35  remapping = remappingTensor->template mutable_data<int>();
36  }
37 
38  const T* input = inputTensor.template data<T>();
39  // TODO(dzhulgakov): if perf becomes an issue consider doing hash table
40  // instead of sorting
41  order_.resize(N);
42  std::iota(order_.begin(), order_.end(), 0);
43  std::sort(order_.begin(), order_.end(), [input](const int x, const int y) {
44  return input[x] < input[y];
45  });
46  int K = N;
47  for (int i = 1; i < N; ++i) {
48  K -= input[order_[i]] == input[order_[i - 1]];
49  }
50  auto* uniqueTensor = Output(UNIQUE, {K}, at::dtype<T>());
51  T* unique = uniqueTensor->template mutable_data<T>();
52  K = 0;
53  T prev = -1;
54  for (int i = 0; i < N; ++i) {
55  if (i == 0 || prev != input[order_[i]]) {
56  prev = unique[K++] = input[order_[i]];
57  }
58  if (remapping) {
59  remapping[order_[i]] = K - 1;
60  }
61  }
62  return true;
63 }
64 
65 REGISTER_CPU_OPERATOR(Unique, UniqueOp<CPUContext>);
66 
67 OPERATOR_SCHEMA(Unique)
68  .NumInputs(1)
69  .NumOutputs(1, 2)
70  .SetDoc(R"DOC(
71 Deduplicates input indices vector and optionally produces reverse remapping.
72 There's no guarantees on the ordering of the output indices.
73 )DOC")
74  .Input(0, "indices", "1D tensor of int32 or int64 indices.")
75  .Output(0, "unique_indices", "1D tensor of deduped entries.")
76  .Output(
77  1,
78  "remapping",
79  "(optional) mapping from `indices` to `unique_indices`. This has the "
80  "same shape as `indices`. Its elements are the indices into "
81  "`unique_indices` such that `Gather(['unique_indices', 'remapping'])` "
82  "yields `indices`.")
83  .TensorInferenceFunction([](const OperatorDef& def,
84  const vector<TensorShape>& in) {
85  std::vector<TensorShape> out(1);
86  out[0].set_data_type(in[0].data_type());
87  CAFFE_ENFORCE_EQ(in[0].dims_size(), 1);
88  if (in[0].dims(0) <= 1) {
89  // This special case is useful in some situation, e.g., when feeding
90  // tensor inference with empty tensor (where the first dim is the batch
91  // size)
92  out[0].add_dims(in[0].dims(0));
93  } else {
94  out[0].set_unknown_shape(true);
95  }
96  if (def.output_size() > 1) {
97  // Remapping has the same shape as the input tensor
98  out.push_back(in[0]);
99  out.back().set_data_type(TensorProto::INT32);
100  }
101  return out;
102  });
103 
104 SHOULD_NOT_DO_GRADIENT(Unique);
105 
106 } // namespace caffe2
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13