Caffe2 - C++ API
A deep learning, cross platform ML framework
top_k.cc
1 
17 #include "caffe2/operators/top_k.h"
18 
19 #include "caffe2/proto/caffe2.pb.h"
20 
21 namespace caffe2 {
22 
23 namespace {
24 
25 template <typename T>
26 struct ValueCmp {
27  bool operator()(
28  const std::pair<T, TIndex>& lhs,
29  const std::pair<T, TIndex>& rhs) {
30  return (
31  lhs.first > rhs.first ||
32  (lhs.first == rhs.first && lhs.second < rhs.second));
33  }
34 };
35 
36 // Define these two names to allow lookup into the 2d tensors like
37 // mytensor(i, j)
38 template <typename T>
39 using EigenMatrixMapRowMajor = Eigen::Map<
40  Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>;
41 
42 template <typename T>
43 using ConstEigenMatrixMapRowMajor = Eigen::Map<
44  const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>;
45 
46 } // namespace
47 
48 template <typename T, class Context>
49 bool TopKOp<T, Context>::RunOnDevice() {
50  auto& input = Input(0);
51  auto* values = Output(0);
52  auto* indices = Output(1);
53  auto* flatten_indices = OutputSize() > 2 ? Output(2) : nullptr;
54 
55  vector<TIndex> in_dims = input.dims();
56  // Linearize input tensor except for last dimension
57  // e.g. [3, 4, 5] -> [12, 5]
58  // [5] -> [5]
59  CAFFE_ENFORCE(
60  in_dims.back() >= k_, "k argment should not be greater than last dim");
61  vector<TIndex> linear_shape = {size_to_dim_(in_dims.size() - 1, in_dims),
62  in_dims[in_dims.size() - 1]};
63  auto input_map = ConstEigenMatrixMapRowMajor<T>(
64  static_cast<const T*>(input.raw_data()),
65  linear_shape[0],
66  linear_shape[1]);
67 
68  // Resize output tensors to be the same shape as the linearized input except
69  // for the last dimension, which will be of size k. E.x. for an input tensor
70  // of shape [3, 4, 5] and k=2, both of these will be shape [3, 4, 2]
71  vector<TIndex> output_linear_shape = {linear_shape[0], k_};
72  values->Resize(output_linear_shape);
73  indices->Resize(output_linear_shape);
74  if (flatten_indices) {
75  flatten_indices->Resize(linear_shape[0] * k_);
76  }
77 
78  // Use Eigen maps to allow indexing into the 2d tensors like values_map(i,j)
79  auto values_map = EigenMatrixMapRowMajor<T>(
80  values->template mutable_data<T>(), linear_shape[0], k_);
81  auto indices_map = EigenMatrixMapRowMajor<TIndex>(
82  indices->template mutable_data<TIndex>(), linear_shape[0], k_);
83  auto* flatten_indices_data = flatten_indices
84  ? flatten_indices->template mutable_data<TIndex>()
85  : nullptr;
86 
87  TIndex flatten_offset = 0;
88  // Sort preserving indices
89  for (TIndex i = 0; i < linear_shape[0]; ++i) {
90  // Build a min-heap, the heap element is pair of (value, idx)
91  // the top of the heap is the smallest value
92  std::priority_queue<
93  std::pair<T, TIndex>,
94  std::vector<std::pair<T, TIndex>>,
95  ValueCmp<T>>
96  PQ;
97 
98  // Maintain the size of heap to be less or equal to k_, so the
99  // heap will hold the k_ largest values
100  for (TIndex j = 0; j < linear_shape[1]; ++j) {
101  const auto value = input_map(i, j);
102  if (PQ.size() < k_ || value > PQ.top().first) {
103  PQ.push(std::make_pair(value, j));
104  }
105  if (PQ.size() > k_) {
106  PQ.pop();
107  }
108  }
109  for (TIndex j = 0; j < k_; ++j) {
110  auto& pqElem = PQ.top();
111  values_map(i, k_ - j - 1) = pqElem.first;
112  indices_map(i, k_ - j - 1) = pqElem.second;
113  if (flatten_indices_data) {
114  flatten_indices_data[k_ - j - 1] = pqElem.second + flatten_offset;
115  }
116  PQ.pop();
117  }
118  if (flatten_indices_data) {
119  flatten_indices_data += k_;
120  }
121  flatten_offset += linear_shape[1];
122  }
123 
124  // Reshape output tensors to [a_1, a_2, ..., a_n, k]
125  auto out_dims = in_dims;
126  out_dims[out_dims.size() - 1] = k_;
127  values->Reshape(out_dims);
128  indices->Reshape(out_dims);
129  return true;
130 }
131 
132 template <typename T, class Context>
133 bool TopKGradientOp<T, Context>::RunOnDevice() {
134  auto& values = Input(0);
135  auto& indices = Input(1);
136  auto& original_input = Input(2);
137  auto* output = Output(0);
138 
139  vector<TIndex> in_dims = values.dims();
140  // Linearize input tensor except for last dimension
141  // e.g. [3, 4, 5] -> [12, 5]
142  // [5] -> [5]
143  vector<TIndex> linear_shape = {size_to_dim_(in_dims.size() - 1, in_dims),
144  in_dims[in_dims.size() - 1]};
145  auto values_map = ConstEigenMatrixMapRowMajor<T>(
146  static_cast<const T*>(values.raw_data()),
147  linear_shape[0],
148  linear_shape[1]);
149  auto indices_map = ConstEigenMatrixMapRowMajor<TIndex>(
150  static_cast<const TIndex*>(indices.raw_data()),
151  linear_shape[0],
152  linear_shape[1]);
153 
154  // Resize output tensors to be as orignial_input size and initialized with 0
155  vector<TIndex> original_dims = original_input.dims();
156  output->Resize(original_dims);
157  T* output_data = output->template mutable_data<T>();
158  memset(output_data, 0, output->nbytes());
159 
160  // Use Eigen maps to allow indexing into the 2d tensors
161  auto output_map = EigenMatrixMapRowMajor<T>(
162  output_data, linear_shape[0], original_dims.back());
163 
164  for (TIndex i = 0; i < linear_shape[0]; ++i) {
165  for (TIndex j = 0; j < linear_shape[1]; ++j) {
166  output_map(i, indices_map(i, j)) = values_map(i, j);
167  }
168  }
169 
170  return true;
171 }
172 
173 REGISTER_CPU_OPERATOR(TopK, TopKOp<float, CPUContext>);
174 REGISTER_CPU_OPERATOR(TopKGradient, TopKGradientOp<float, CPUContext>);
175 
176 OPERATOR_SCHEMA(TopK)
177  .NumInputs(1)
178  .NumOutputs(2, 3)
179  .TensorInferenceFunction([](const OperatorDef& def,
180  const vector<TensorShape>& in) {
181  vector<TensorShape> out = {in[0], in[0]};
182  ArgumentHelper helper(def);
183  auto k = helper.GetSingleArgument("k", -1);
184  auto dims_size = in[0].dims_size();
185  out[0].set_dims(dims_size - 1, k);
186  out[1].set_dims(dims_size - 1, k);
187  out[1].set_data_type(TensorProto_DataType_INT32);
188  if (def.output_size() > 2) {
189  TensorShape flatten_indices_shape;
190  flatten_indices_shape.set_data_type(TensorProto_DataType_INT32);
191  flatten_indices_shape.add_dims(
192  std::accumulate(
193  in[0].dims().begin(),
194  in[0].dims().end() - 1,
195  1,
196  std::multiplies<long>()) *
197  k);
198  out.push_back(flatten_indices_shape);
199  }
200  return out;
201  })
202  .SetDoc(R"DOC(
203 Retrieve the top-K elements for the last dimension. Given an input tensor of
204 shape [a_1, a_2, ..., a_n, r] and integer argument k, return two outputs:
205 -Value tensor of shape [a_1, a_2, ..., a_n, k] which contains the values of
206  the top k elements along the last dimension
207 -Index tensor of shape [a_1, a_2, ..., a_n, k] which contains the indices
208  of the top k elements (original indices from the input tensor).
209 
210 Given two equivalent values, this operator uses the indices along the last dim-
211 ension as a tiebreaker. That is, the element with the lower index will appear
212 first.
213  )DOC")
214  .Input(0, "X", "Tensor of shape [a_1, a_2, ..., a_n, r]")
215  .Output(
216  0,
217  "Values",
218  "Tensor of shape [a_1, a_2, ..., a_n, k] containing"
219  " top K values from the input tensor")
220  .Output(
221  1,
222  "Indices",
223  "Tensor of shape [a_1, a_2, ..., a_n, k] containing"
224  " the corresponding input tensor indices for the top K values.")
225  .Output(
226  2,
227  "Flatten indices",
228  "Tensor of shape [a_1 * a_2 * ... * a_n * k] containing the indices "
229  "into the flatten input")
230  .Arg("k", "Number of top elements to retrieve");
231 
232 OPERATOR_SCHEMA(TopKGradient).NumInputs(3).NumOutputs(1);
233 
235  using GradientMakerBase::GradientMakerBase;
236  vector<OperatorDef> GetGradientDefs() override {
237  return SingleGradientDef(
238  "TopKGradient",
239  "",
240  vector<string>{GO(0), O(1), I(0)},
241  vector<string>{GI(0)});
242  }
243 };
244 
245 REGISTER_GRADIENT(TopK, GetTopKGradient);
246 
247 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.