Caffe2 - C++ API
A deep learning, cross platform ML framework
find_duplicate_elements_op.h
1 
17 #ifndef CAFFE2_OPERATORS_FIND_DUPLICATE_ELEMENTS_OP_H
18 #define CAFFE2_OPERATORS_FIND_DUPLICATE_ELEMENTS_OP_H
19 
20 #include <unordered_map>
21 #include <vector>
22 
23 #include "caffe2/core/context.h"
24 #include "caffe2/core/operator.h"
25 #include "caffe2/core/tensor.h"
26 
27 namespace caffe2 {
28 
29 template <class Context>
30 class FindDuplicateElementsOp final : public Operator<Context> {
31  public:
32  USE_OPERATOR_CONTEXT_FUNCTIONS;
33  USE_SIMPLE_CTOR_DTOR(FindDuplicateElementsOp);
34  USE_DISPATCH_HELPER;
35 
36  bool RunOnDevice() override {
38  call(this, Input(0));
39  }
40 
41  template <typename T>
42  bool DoRunWithType() {
43  const auto& data = Input(0);
44  CAFFE_ENFORCE(data.ndim() == 1, "data should be 1-D.");
45 
46  const auto* data_ptr = data.template data<T>();
47  std::unordered_map<T, int64_t> dict;
48  std::vector<int64_t> dupIndices;
49  // i is the index of unique elements, j is the index of all elements
50  for (int64_t i = 0, j = 0; j < data.dims()[0]; ++i, ++j) {
51  bool retVal = dict.insert({data_ptr[j], i}).second;
52  if (!retVal) {
53  --i;
54  dupIndices.push_back(j);
55  }
56  }
57 
58  const auto dupSize = dupIndices.size();
59  auto* output = Output(0);
60  output->Resize(dupSize);
61  auto* out_ptr = output->template mutable_data<int64_t>();
62  for (int64_t i = 0; i < dupSize; ++i) {
63  out_ptr[i] = dupIndices[i];
64  }
65 
66  return true;
67  }
68 };
69 
70 } // namespace caffe2
71 
72 #endif // CAFFE2_OPERATORS_FIND_DUPLICATE_ELEMENTS_OP_H
Copyright (c) 2016-present, Facebook, Inc.