Caffe2 - C++ API
A deep learning, cross platform ML framework
apmeter_op.cc
1 #include "caffe2/operators/apmeter_op.h"
2 
3 namespace caffe2 {
4 
5 template <>
6 void APMeterOp<float, CPUContext>::BufferPredictions(
7  const float* XData,
8  const int* labelData,
9  int N,
10  int D) {
11  if (buffers_.empty()) {
12  // Initialize the buffer
13  buffers_.resize(D, std::vector<BufferDataType>(buffer_size_));
14  }
15  DCHECK_EQ(buffers_.size(), D);
16 
17  // Fill atmose buffer_size_ data at a time, so truncate the input if needed
18  if (N > buffer_size_) {
19  XData = XData + (N - buffer_size_) * D;
20  labelData = labelData + (N - buffer_size_) * D;
21  N = buffer_size_;
22  }
23 
24  // Reclaim space if not enough space in the buffer to hold new data
25  int space_to_reclaim = buffer_used_ + N - buffer_size_;
26  if (space_to_reclaim > 0) {
27  for (auto& buffer : buffers_) {
28  std::rotate(
29  buffer.begin(), buffer.begin() + space_to_reclaim, buffer.end());
30  }
31  buffer_used_ -= space_to_reclaim;
32  }
33 
34  // Fill the buffer
35  for (int i = 0; i < D; i++) {
36  for (int j = 0; j < N; j++) {
37  buffers_[i][buffer_used_ + j].first = XData[j * D + i];
38  buffers_[i][buffer_used_ + j].second = labelData[j * D + i];
39  }
40  }
41 
42  buffer_used_ += N;
43 }
44 
45 template <>
46 bool APMeterOp<float, CPUContext>::RunOnDevice() {
47  auto& X = Input(PREDICTION);
48  auto& label = Input(LABEL);
49 
50  // Check dimensions
51  DCHECK_EQ(X.dim(), 2);
52  int N = X.dim32(0);
53  int D = X.dim32(1);
54  DCHECK_EQ(label.dim(), 2);
55  DCHECK_EQ(label.dim32(0), N);
56  DCHECK_EQ(label.dim32(1), D);
57  auto* Y = Output(0, {D}, at::dtype<float>());
58 
59  const auto* Xdata = X.data<float>();
60  const auto* labelData = label.data<int>();
61  auto* Ydata = Y->template mutable_data<float>();
62 
63  BufferPredictions(Xdata, labelData, N, D);
64 
65  // Calculate AP for each class
66  for (int i = 0; i < D; i++) {
67  auto& buffer = buffers_[i];
68  // Sort predictions by score
69  std::stable_sort(
70  buffer.begin(),
71  buffer.begin() + buffer_used_,
72  [](const BufferDataType& p1, const BufferDataType& p2) {
73  return p1.first > p2.first;
74  });
75  // Calculate cumulative precision for each sample
76  float tp_sum = 0.0;
77  float precision_sum = 0.0;
78  int ntruth = 0;
79  for (int j = 0; j < buffer_used_; j++) {
80  tp_sum += buffer[j].second;
81  if (buffer[j].second == 1) {
82  ntruth += 1;
83  precision_sum += tp_sum / (j + 1);
84  }
85  }
86 
87  // Calculate AP
88  Ydata[i] = precision_sum / std::max(1, ntruth);
89  }
90 
91  return true;
92 }
93 
94 namespace {
95 REGISTER_CPU_OPERATOR(APMeter, APMeterOp<float, CPUContext>);
96 
97 OPERATOR_SCHEMA(APMeter)
98  .NumInputs(2)
99  .NumOutputs(1)
100  .ScalarType(TensorProto::FLOAT)
101  .SetDoc(R"DOC(
102 APMeter computes Average Precision for binary or multi-class classification.
103 It takes two inputs: prediction scores P of size (n_samples x n_classes), and
104 true labels Y of size (n_samples x n_classes). It returns a single float number
105 per class for the average precision of that class.
106 )DOC")
107  .Arg(
108  "buffer_size",
109  "(int32_t) indicates how many predictions should the op buffer. "
110  "defaults to 1000")
111  .Input(
112  0,
113  "predictions",
114  "2-D tensor (Tensor<float>) of size (num_samples x"
115  "num_classes) containing prediction scores")
116  .Input(
117  1,
118  "labels",
119  "2-D tensor (Tensor<float>) of size (num_samples) "
120  "containing true labels for each sample")
121  .Output(
122  0,
123  "AP",
124  "1-D tensor (Tensor<float>) of size num_classes containing "
125  "average precision for each class");
126 
127 SHOULD_NOT_DO_GRADIENT(APMeter);
128 
129 } // namespace
130 } // namespace caffe2
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13
Definition: static.cpp:70