Simplify the top-k argmax layer using std::sort

diff --git a/include/caffe/common_layers.hpp b/include/caffe/common_layers.hpp
index 395ea7d..26bfb5d 100644 (file)
--- a/include/caffe/common_layers.hpp
+++ b/include/caffe/common_layers.hpp
@@ -17,6 +17,12 @@
 
 namespace caffe {
 
+template<typename Dtype>
+bool int_Dtype_pair_greater(std::pair<int, Dtype> a,
+                            std::pair<int, Dtype> b) {
+  return a.second > b.second || (a.second == b.second && a.first > b.first);
+}
+
 /* ArgmaxLayer
   Compute the index of the max value across all (channels x height x width).
   [In the future, can take specific dimension.]

diff --git a/src/caffe/layers/argmax_layer.cpp b/src/caffe/layers/argmax_layer.cpp
index 1fe2a89..00c17f1 100644 (file)
--- a/src/caffe/layers/argmax_layer.cpp
+++ b/src/caffe/layers/argmax_layer.cpp
@@ -1,5 +1,6 @@
 // Copyright 2014 BVLC and contributors.
 
+#include <algorithm>
 #include <cfloat>
 #include <queue>
 #include <vector>
@@ -28,50 +29,27 @@ void ArgMaxLayer<Dtype>::SetUp(const vector<Blob<Dtype>*>& bottom,
 }
 
 template <typename Dtype>
-class IDAndValueComparator {
- public:
-  bool operator() (const std::pair<size_t, Dtype>& lhs,
-                   const std::pair<size_t, Dtype>& rhs) const {
-    return lhs.second < rhs.second || (lhs.second == rhs.second &&
-        lhs.first < rhs.first);
-  }
-};
-
-template <typename Dtype>
 Dtype ArgMaxLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
     vector<Blob<Dtype>*>* top) {
   const Dtype* bottom_data = bottom[0]->cpu_data();
   Dtype* top_data = (*top)[0]->mutable_cpu_data();
   int num = bottom[0]->num();
   int dim = bottom[0]->count() / bottom[0]->num();
-  Dtype value;
   for (int i = 0; i < num; ++i) {
-    std::priority_queue<std::pair<size_t, Dtype>,
-        std::vector<std::pair<size_t, Dtype> >, IDAndValueComparator<Dtype> >
-        top_k_results;
+    std::vector<std::pair<int, Dtype> > bottom_data_vector;
     for (int j = 0; j < dim; ++j) {
-      value = -(bottom_data[i * dim + j]);
-      if (top_k_results.size() >= top_k_) {
-        if (value < top_k_results.top().second) {
-          top_k_results.pop();
-          top_k_results.push(std::make_pair(j, value));
-        }
-      } else {
-        top_k_results.push(std::make_pair(j, value));
-      }
+      bottom_data_vector.push_back(std::make_pair(j, bottom_data[i * dim + j]));
     }
+    std::sort(bottom_data_vector.begin(), bottom_data_vector.end(),
+              int_Dtype_pair_greater<Dtype>);
     if (out_max_val_) {
       for (int j = 0; j < top_k_; ++j) {
-        top_data[i * 2 * top_k_ + (top_k_ - 1 - j) * 2] =
-            top_k_results.top().first;
-        top_data[i * 2 * top_k_ + (top_k_ - 1 - j) * 2 + 1] =
-            -(top_k_results.top().second);
-        top_k_results.pop();
+        top_data[(*top)[0]->offset(i, 0, j)] = bottom_data_vector[j].first;
+        top_data[(*top)[0]->offset(i, 1, j)] = bottom_data_vector[j].second;
       }
     } else {
       for (int j = 0; j < top_k_; ++j) {
-        top_data[i * top_k_ + (top_k_ - 1 - j)] = top_k_results.top().first;
-        top_k_results.pop();
+        top_data[(*top)[0]->offset(i, 0, j)] = bottom_data_vector[j].first;
       }
     }
   }

diff --git a/src/caffe/test/test_argmax_layer.cpp b/src/caffe/test/test_argmax_layer.cpp
index 11b6712..5286a0a 100644 (file)
--- a/src/caffe/test/test_argmax_layer.cpp
+++ b/src/caffe/test/test_argmax_layer.cpp
@@ -1,6 +1,5 @@
 // Copyright 2014 BVLC and contributors.
 
-#include <algorithm>
 #include <utility>
 #include <vector>
 
@@ -109,12 +108,6 @@ TYPED_TEST(ArgMaxLayerTest, TestCPUMaxVal) {
   }
 }
 
-template<typename Dtype>
-bool int_Dtype_pair_greater(std::pair<int, Dtype> a,
-                            std::pair<int, Dtype> b) {
-  return a.second > b.second || (a.second == b.second && a.first > b.first);
-}
-
 TYPED_TEST(ArgMaxLayerTest, TestCPUTopK) {
   LayerParameter layer_param;
   Caffe::set_mode(Caffe::CPU);
@@ -124,27 +117,23 @@ TYPED_TEST(ArgMaxLayerTest, TestCPUTopK) {
   layer.SetUp(this->blob_bottom_vec_, &(this->blob_top_vec_));
   layer.Forward(this->blob_bottom_vec_, &(this->blob_top_vec_));
   // Now, check values
-  const TypeParam* bottom_data = this->blob_bottom_->cpu_data();
-  const TypeParam* top_data = this->blob_top_->cpu_data();
   int max_ind;
   TypeParam max_val;
   int num = this->blob_bottom_->num();
   int dim = this->blob_bottom_->count() / num;
   for (int i = 0; i < num; ++i) {
-    EXPECT_GE(top_data[i], 0);
-    EXPECT_LE(top_data[i], dim);
-    max_ind = top_data[i * this->top_k_];
-    max_val = bottom_data[i * dim + max_ind];
-    EXPECT_EQ(bottom_data[i * dim + max_ind], max_val);
-    std::vector<std::pair<int, TypeParam> > bottom_data_vector;
-    for (int j = 0; j < dim; ++j) {
-      EXPECT_LE(bottom_data[i * dim + j], max_val);
-      bottom_data_vector.push_back(std::make_pair(j, bottom_data[i * dim + j]));
-    }
-    std::sort(bottom_data_vector.begin(), bottom_data_vector.end(),
-              int_Dtype_pair_greater<TypeParam>);
+    EXPECT_GE(this->blob_top_->data_at(i, 0, 0, 0), 0);
+    EXPECT_LE(this->blob_top_->data_at(i, 0, 0, 0), dim);
     for (int j = 0; j < this->top_k_; ++j) {
-      EXPECT_EQ(bottom_data_vector[j].first, top_data[i * this->top_k_ + j]);
+      max_ind = this->blob_top_->data_at(i, 0, j, 0);
+      max_val = this->blob_bottom_->data_at(i, max_ind, 0, 0);
+      int count = 0;
+      for (int k = 0; k < dim; ++k) {
+        if (this->blob_bottom_->data_at(i, k, 0, 0) > max_val) {
+          ++count;
+        }
+      }
+      EXPECT_EQ(j, count);
     }
   }
 }
@@ -159,30 +148,24 @@ TYPED_TEST(ArgMaxLayerTest, TestCPUMaxValTopK) {
   layer.SetUp(this->blob_bottom_vec_, &(this->blob_top_vec_));
   layer.Forward(this->blob_bottom_vec_, &(this->blob_top_vec_));
   // Now, check values
-  const TypeParam* bottom_data = this->blob_bottom_->cpu_data();
-  const TypeParam* top_data = this->blob_top_->cpu_data();
   int max_ind;
   TypeParam max_val;
   int num = this->blob_bottom_->num();
   int dim = this->blob_bottom_->count() / num;
   for (int i = 0; i < num; ++i) {
-    EXPECT_GE(top_data[i], 0);
-    EXPECT_LE(top_data[i], dim);
-    max_ind = top_data[i * 2 * this->top_k_];
-    max_val = top_data[i * 2 * this->top_k_ + 1];
-    EXPECT_EQ(bottom_data[i * dim + max_ind], max_val);
-    std::vector<std::pair<int, TypeParam> > bottom_data_vector;
-    for (int j = 0; j < dim; ++j) {
-      EXPECT_LE(bottom_data[i * dim + j], max_val);
-      bottom_data_vector.push_back(std::make_pair(j, bottom_data[i * dim + j]));
-    }
-    std::sort(bottom_data_vector.begin(), bottom_data_vector.end(),
-              int_Dtype_pair_greater<TypeParam>);
+    EXPECT_GE(this->blob_top_->data_at(i, 0, 0, 0), 0);
+    EXPECT_LE(this->blob_top_->data_at(i, 0, 0, 0), dim);
     for (int j = 0; j < this->top_k_; ++j) {
-      EXPECT_EQ(bottom_data_vector[j].first,
-                top_data[i * 2 * this->top_k_ + 2 * j]);
-      EXPECT_EQ(bottom_data_vector[j].second,
-                top_data[i * 2 * this->top_k_ + 2 * j + 1]);
+      max_ind = this->blob_top_->data_at(i, 0, j, 0);
+      max_val = this->blob_top_->data_at(i, 1, j, 0);
+      EXPECT_EQ(this->blob_bottom_->data_at(i, max_ind, 0, 0), max_val);
+      int count = 0;
+      for (int k = 0; k < dim; ++k) {
+        if (this->blob_bottom_->data_at(i, k, 0, 0) > max_val) {
+          ++count;
+        }
+      }
+      EXPECT_EQ(j, count);
     }
   }
 }