Merge pull request #2076 from jeffdonahue/accuracy-layer-fixes

Fixup AccuracyLayer like SoftmaxLossLayer in #1970

4 files changed, 154 insertions, 28 deletions

diff --git a/include/caffe/loss_layers.hpp b/include/caffe/loss_layers.hpp
index 62d6df71..d3eecd2e 100644
--- a/include/caffe/loss_layers.hpp
+++ b/include/caffe/loss_layers.hpp
@@ -78,7 +78,14 @@ class AccuracyLayer : public Layer<Dtype> {
     }
   }
 
+  int label_axis_, outer_num_, inner_num_;
+
   int top_k_;
+
+  /// Whether to ignore instances with a certain label.
+  bool has_ignore_label_;
+  /// The label indicating that an instance should be ignored.
+  int ignore_label_;
 };
 
 /**
diff --git a/src/caffe/layers/accuracy_layer.cpp b/src/caffe/layers/accuracy_layer.cpp
index 186f9f86..90aad675 100644
--- a/src/caffe/layers/accuracy_layer.cpp
+++ b/src/caffe/layers/accuracy_layer.cpp
@@ -14,6 +14,12 @@ template <typename Dtype>
 void AccuracyLayer<Dtype>::LayerSetUp(
   const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
   top_k_ = this->layer_param_.accuracy_param().top_k();
+
+  has_ignore_label_ =
+    this->layer_param_.accuracy_param().has_ignore_label();
+  if (has_ignore_label_) {
+    ignore_label_ = this->layer_param_.accuracy_param().ignore_label();
+  }
 }
 
 template <typename Dtype>
@@ -21,11 +27,15 @@ void AccuracyLayer<Dtype>::Reshape(
   const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
   CHECK_LE(top_k_, bottom[0]->count() / bottom[1]->count())
       << "top_k must be less than or equal to the number of classes.";
-  CHECK_GE(bottom[0]->num_axes(), bottom[1]->num_axes());
-  for (int i = 0; i < bottom[1]->num_axes(); ++i) {
-    CHECK_LE(bottom[0]->shape(i), bottom[1]->shape(i))
-        << "Dimension mismatch between predictions and label.";
-  }
+  label_axis_ =
+      bottom[0]->CanonicalAxisIndex(this->layer_param_.accuracy_param().axis());
+  outer_num_ = bottom[0]->count(0, label_axis_);
+  inner_num_ = bottom[0]->count(label_axis_ + 1);
+  CHECK_EQ(outer_num_ * inner_num_, bottom[1]->count())
+      << "Number of labels must match number of predictions; "
+      << "e.g., if label axis == 1 and prediction shape is (N, C, H, W), "
+      << "label count (number of labels) must be N*H*W, "
+      << "with integer values in {0, 1, ..., C-1}.";
   vector<int> top_shape(0);  // Accuracy is a scalar; 0 axes.
   top[0]->Reshape(top_shape);
 }
@@ -36,31 +46,42 @@ void AccuracyLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
   Dtype accuracy = 0;
   const Dtype* bottom_data = bottom[0]->cpu_data();
   const Dtype* bottom_label = bottom[1]->cpu_data();
-  int num = bottom[0]->count(0, bottom[1]->num_axes());
-  int dim = bottom[0]->count() / num;
+  const int dim = bottom[0]->count() / outer_num_;
+  const int num_labels = bottom[0]->shape(label_axis_);
   vector<Dtype> maxval(top_k_+1);
   vector<int> max_id(top_k_+1);
-  for (int i = 0; i < num; ++i) {
-    // Top-k accuracy
-    std::vector<std::pair<Dtype, int> > bottom_data_vector;
-    for (int j = 0; j < dim; ++j) {
-      bottom_data_vector.push_back(
-          std::make_pair(bottom_data[i * dim + j], j));
-    }
-    std::partial_sort(
-        bottom_data_vector.begin(), bottom_data_vector.begin() + top_k_,
-        bottom_data_vector.end(), std::greater<std::pair<Dtype, int> >());
-    // check if true label is in top k predictions
-    for (int k = 0; k < top_k_; k++) {
-      if (bottom_data_vector[k].second == static_cast<int>(bottom_label[i])) {
-        ++accuracy;
-        break;
+  int count = 0;
+  for (int i = 0; i < outer_num_; ++i) {
+    for (int j = 0; j < inner_num_; ++j) {
+      const int label_value =
+          static_cast<int>(bottom_label[i * inner_num_ + j]);
+      if (has_ignore_label_ && label_value == ignore_label_) {
+        continue;
+      }
+      DCHECK_GE(label_value, 0);
+      DCHECK_LT(label_value, num_labels);
+      // Top-k accuracy
+      std::vector<std::pair<Dtype, int> > bottom_data_vector;
+      for (int k = 0; k < num_labels; ++k) {
+        bottom_data_vector.push_back(std::make_pair(
+            bottom_data[i * dim + k * inner_num_ + j], k));
+      }
+      std::partial_sort(
+          bottom_data_vector.begin(), bottom_data_vector.begin() + top_k_,
+          bottom_data_vector.end(), std::greater<std::pair<Dtype, int> >());
+      // check if true label is in top k predictions
+      for (int k = 0; k < top_k_; k++) {
+        if (bottom_data_vector[k].second == label_value) {
+          ++accuracy;
+          break;
+        }
       }
+      ++count;
     }
   }
 
   // LOG(INFO) << "Accuracy: " << accuracy;
-  top[0]->mutable_cpu_data()[0] = accuracy / num;
+  top[0]->mutable_cpu_data()[0] = accuracy / count;
   // Accuracy layer should not be used as a loss function.
 }
 
diff --git a/src/caffe/proto/caffe.proto b/src/caffe/proto/caffe.proto
index 3b479466..e523efa5 100644
--- a/src/caffe/proto/caffe.proto
+++ b/src/caffe/proto/caffe.proto
@@ -367,6 +367,16 @@ message AccuracyParameter {
   // the top k scoring classes.  By default, only compare to the top scoring
   // class (i.e. argmax).
   optional uint32 top_k = 1 [default = 1];
+
+  // The "label" axis of the prediction blob, whose argmax corresponds to the
+  // predicted label -- may be negative to index from the end (e.g., -1 for the
+  // last axis).  For example, if axis == 1 and the predictions are
+  // (N x C x H x W), the label blob is expected to contain N*H*W ground truth
+  // labels with integer values in {0, 1, ..., C-1}.
+  optional int32 axis = 2 [default = 1];
+
+  // If specified, ignore instances with the given label.
+  optional int32 ignore_label = 3;
 }
 
 // Message that stores parameters used by ArgMaxLayer
diff --git a/src/caffe/test/test_accuracy_layer.cpp b/src/caffe/test/test_accuracy_layer.cpp
index 1c58b767..6cbf51df 100644
--- a/src/caffe/test/test_accuracy_layer.cpp
+++ b/src/caffe/test/test_accuracy_layer.cpp
@@ -29,6 +29,14 @@ class AccuracyLayerTest : public ::testing::Test {
     blob_bottom_data_->Reshape(shape);
     shape.resize(1);
     blob_bottom_label_->Reshape(shape);
+    FillBottoms();
+
+    blob_bottom_vec_.push_back(blob_bottom_data_);
+    blob_bottom_vec_.push_back(blob_bottom_label_);
+    blob_top_vec_.push_back(blob_top_);
+  }
+
+  virtual void FillBottoms() {
     // fill the probability values
     FillerParameter filler_param;
     GaussianFiller<Dtype> filler(filler_param);
@@ -39,14 +47,11 @@ class AccuracyLayerTest : public ::testing::Test {
     caffe::rng_t* prefetch_rng =
           static_cast<caffe::rng_t*>(rng->generator());
     Dtype* label_data = blob_bottom_label_->mutable_cpu_data();
-    for (int i = 0; i < 100; ++i) {
+    for (int i = 0; i < blob_bottom_label_->count(); ++i) {
       label_data[i] = (*prefetch_rng)() % 10;
     }
-
-    blob_bottom_vec_.push_back(blob_bottom_data_);
-    blob_bottom_vec_.push_back(blob_bottom_label_);
-    blob_top_vec_.push_back(blob_top_);
   }
+
   virtual ~AccuracyLayerTest() {
     delete blob_bottom_data_;
     delete blob_bottom_label_;
@@ -112,6 +117,89 @@ TYPED_TEST(AccuracyLayerTest, TestForwardCPU) {
               num_correct_labels / 100.0, 1e-4);
 }
 
+TYPED_TEST(AccuracyLayerTest, TestForwardWithSpatialAxes) {
+  Caffe::set_mode(Caffe::CPU);
+  this->blob_bottom_data_->Reshape(2, 10, 4, 5);
+  vector<int> label_shape(3);
+  label_shape[0] = 2; label_shape[1] = 4; label_shape[2] = 5;
+  this->blob_bottom_label_->Reshape(label_shape);
+  this->FillBottoms();
+  LayerParameter layer_param;
+  layer_param.mutable_accuracy_param()->set_axis(1);
+  AccuracyLayer<TypeParam> layer(layer_param);
+  layer.SetUp(this->blob_bottom_vec_, this->blob_top_vec_);
+  layer.Forward(this->blob_bottom_vec_, this->blob_top_vec_);
+
+  TypeParam max_value;
+  const int num_labels = this->blob_bottom_label_->count();
+  int max_id;
+  int num_correct_labels = 0;
+  vector<int> label_offset(3);
+  for (int n = 0; n < this->blob_bottom_data_->num(); ++n) {
+    for (int h = 0; h < this->blob_bottom_data_->height(); ++h) {
+      for (int w = 0; w < this->blob_bottom_data_->width(); ++w) {
+        max_value = -FLT_MAX;
+        max_id = 0;
+        for (int c = 0; c < this->blob_bottom_data_->channels(); ++c) {
+          const TypeParam pred_value =
+              this->blob_bottom_data_->data_at(n, c, h, w);
+          if (pred_value > max_value) {
+            max_value = pred_value;
+            max_id = c;
+          }
+        }
+        label_offset[0] = n; label_offset[1] = h; label_offset[2] = w;
+        const int correct_label =
+            static_cast<int>(this->blob_bottom_label_->data_at(label_offset));
+        if (max_id == correct_label) {
+          ++num_correct_labels;
+        }
+      }
+    }
+  }
+  EXPECT_NEAR(this->blob_top_->data_at(0, 0, 0, 0),
+              num_correct_labels / TypeParam(num_labels), 1e-4);
+}
+
+TYPED_TEST(AccuracyLayerTest, TestForwardIgnoreLabel) {
+  Caffe::set_mode(Caffe::CPU);
+  LayerParameter layer_param;
+  const TypeParam kIgnoreLabelValue = -1;
+  layer_param.mutable_accuracy_param()->set_ignore_label(kIgnoreLabelValue);
+  AccuracyLayer<TypeParam> layer(layer_param);
+  // Manually set some labels to the ignore label value (-1).
+  this->blob_bottom_label_->mutable_cpu_data()[2] = kIgnoreLabelValue;
+  this->blob_bottom_label_->mutable_cpu_data()[5] = kIgnoreLabelValue;
+  this->blob_bottom_label_->mutable_cpu_data()[32] = kIgnoreLabelValue;
+  layer.SetUp(this->blob_bottom_vec_, this->blob_top_vec_);
+  layer.Forward(this->blob_bottom_vec_, this->blob_top_vec_);
+
+  TypeParam max_value;
+  int max_id;
+  int num_correct_labels = 0;
+  int count = 0;
+  for (int i = 0; i < 100; ++i) {
+    if (kIgnoreLabelValue == this->blob_bottom_label_->data_at(i, 0, 0, 0)) {
+      continue;
+    }
+    ++count;
+    max_value = -FLT_MAX;
+    max_id = 0;
+    for (int j = 0; j < 10; ++j) {
+      if (this->blob_bottom_data_->data_at(i, j, 0, 0) > max_value) {
+        max_value = this->blob_bottom_data_->data_at(i, j, 0, 0);
+        max_id = j;
+      }
+    }
+    if (max_id == this->blob_bottom_label_->data_at(i, 0, 0, 0)) {
+      ++num_correct_labels;
+    }
+  }
+  EXPECT_EQ(count, 97);  // We set 3 out of 100 labels to kIgnoreLabelValue.
+  EXPECT_NEAR(this->blob_top_->data_at(0, 0, 0, 0),
+              num_correct_labels / TypeParam(count), 1e-4);
+}
+
 TYPED_TEST(AccuracyLayerTest, TestForwardCPUTopK) {
   LayerParameter layer_param;
   AccuracyParameter* accuracy_param = layer_param.mutable_accuracy_param();