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