// Copyright 2014 BVLC and contributors.
+#include <algorithm>
#include <cfloat>
#include <queue>
#include <vector>
}
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;
}
}
}
// Copyright 2014 BVLC and contributors.
-#include <algorithm>
#include <utility>
#include <vector>
}
}
-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);
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);
}
}
}
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);
}
}
}