Caffe2 - C++ API
A deep learning, cross platform ML framework
remove_data_blocks_op.h
1 
17 #ifndef CAFFE2_OPERATORS_REMOVE_DATA_BLOCKS_OP_H_
18 #define CAFFE2_OPERATORS_REMOVE_DATA_BLOCKS_OP_H_
19 
20 #include <algorithm>
21 #include <vector>
22 
23 #include "caffe2/core/context.h"
24 #include "caffe2/core/operator.h"
25 
26 namespace caffe2 {
27 
28 template <class Context>
29 class RemoveDataBlocksOp final : public Operator<Context> {
30  public:
31  USE_OPERATOR_CONTEXT_FUNCTIONS;
32  USE_SIMPLE_CTOR_DTOR(RemoveDataBlocksOp);
33  USE_DISPATCH_HELPER;
34 
35  bool RunOnDevice() override {
36  if (Input(INDICES).dims()[0] == 0) {
37  Output(0)->CopyFrom(Input(0));
38  return true;
39  } else {
40  return DispatchHelper<TensorTypes<int, long>>::call(this, Input(INDICES));
41  }
42  }
43 
44  template <typename T>
45  bool DoRunWithType() {
46  const auto& data = Input(DATA);
47  const auto& indices = Input(INDICES);
48  CAFFE_ENFORCE(data.ndim() > 0, "DATA should be at leat 1-D.");
49  CAFFE_ENFORCE(indices.ndim() == 1, "INDICES should be 1-D.");
50 
51  const auto outer_size = data.dims()[0];
52  const auto block_size = data.size_from_dim(1);
53  const auto block_size_bytes = block_size * data.meta().itemsize();
54  auto indices_size = indices.dims()[0];
55  const char* data_ptr = (char*)data.raw_data();
56  const auto* ind_ptr = indices.template data<T>();
57 
58  std::vector<T> ind_vec;
59  for (int64_t i = 0; i < indices_size; i++) {
60  ind_vec.push_back(ind_ptr[i]);
61  }
62  std::sort(ind_vec.begin(), ind_vec.end());
63  CAFFE_ENFORCE(ind_vec[0] >= 0, "The min index should be larger than zero.");
64  CAFFE_ENFORCE(
65  ind_vec[indices_size - 1] < outer_size,
66  "The max index should be smaller than the data outer size.");
67  // removes duplicate indices
68  ind_vec.erase(std::unique(ind_vec.begin(), ind_vec.end()), ind_vec.end());
69  indices_size = ind_vec.size();
70 
71  auto* output = Output(0);
72  auto shape = data.dims();
73  shape[0] -= indices_size;
74  output->Resize(shape);
75  char* out_ptr = (char*)output->raw_mutable_data(data.meta());
76 
77  ind_vec.insert(ind_vec.begin(), -1);
78  int64_t ind_vec_size = ind_vec.size();
79  for (auto i = 0; i < ind_vec_size; i++) {
80  int64_t interval_start = ind_vec[i] + 1;
81  int64_t interval_end =
82  (i == ind_vec_size - 1) ? outer_size : ind_vec[i + 1];
83  auto num_items = interval_end - interval_start;
84  context_.template CopyItems<Context, Context>(
85  data.meta(),
86  num_items * block_size,
87  data_ptr + block_size_bytes * interval_start,
88  out_ptr);
89  out_ptr += block_size_bytes * num_items;
90  }
91 
92  return true;
93  }
94 
95  private:
96  INPUT_TAGS(DATA, INDICES);
97 };
98 
99 } // namespace caffe2
100 
101 #endif // CAFFE2_OPERATORS_REMOVE_DATA_BLOCKS_OP_H_
Copyright (c) 2016-present, Facebook, Inc.