fix really stupid bug in flatten layer (and add test that shows the

failure case; not sure why CheckGradientExhaustive didn't catch it)
skyuuka · Dec 5, 2013 · b906327 · b906327
1 parent 8c96ac2
commit b906327
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 8 deletions.
diff --git a/include/caffe/vision_layers.hpp b/include/caffe/vision_layers.hpp
@@ -106,7 +106,7 @@ class FlattenLayer : public Layer<Dtype> {
       const bool propagate_down, vector<Blob<Dtype>*>* bottom);
   virtual Dtype Backward_gpu(const vector<Blob<Dtype>*>& top,
       const bool propagate_down, vector<Blob<Dtype>*>* bottom);
-  int channels_out_;
+  int count_;
 };
 
 

diff --git a/src/caffe/layers/flatten_layer.cpp b/src/caffe/layers/flatten_layer.cpp
@@ -13,34 +13,36 @@ void FlattenLayer<Dtype>::SetUp(const vector<Blob<Dtype>*>& bottom,
       vector<Blob<Dtype>*>* top) {
   CHECK_EQ(bottom.size(), 1) << "Flatten Layer takes a single blob as input.";
   CHECK_EQ(top->size(), 1) << "Flatten Layer takes a single blob as output.";
-  channels_out_ = bottom[0]->channels() * bottom[0]->height()
+  int channels_out = bottom[0]->channels() * bottom[0]->height()
       * bottom[0]->width();
-  (*top)[0]->Reshape(bottom[0]->num(), channels_out_, 1, 1);
+  (*top)[0]->Reshape(bottom[0]->num(), channels_out, 1, 1);
+  count_ = bottom[0]->num() * channels_out;
+  CHECK_EQ(count_, bottom[0]->count());
+  CHECK_EQ(count_, (*top)[0]->count());
 };
 
 template <typename Dtype>
 void FlattenLayer<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();
-  caffe_copy(channels_out_, bottom_data, top_data);
+  caffe_copy(count_, bottom_data, top_data);
 }
 
 template <typename Dtype>
 void FlattenLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype>*>& bottom,
       vector<Blob<Dtype>*>* top) {
   const Dtype* bottom_data = bottom[0]->gpu_data();
   Dtype* top_data = (*top)[0]->mutable_gpu_data();
-  caffe_gpu_copy(channels_out_, bottom_data, top_data);
+  caffe_gpu_copy(count_, bottom_data, top_data);
 }
 
 template <typename Dtype>
 Dtype FlattenLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
       const bool propagate_down, vector<Blob<Dtype>*>* bottom) {
   const Dtype* top_diff = top[0]->cpu_diff();
   Dtype* bottom_diff = (*bottom)[0]->mutable_cpu_diff();
-  caffe_copy(channels_out_, top_diff, bottom_diff);
+  caffe_copy(count_, top_diff, bottom_diff);
 }
 
 
@@ -49,7 +51,7 @@ Dtype FlattenLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype>*>& top,
       const bool propagate_down, vector<Blob<Dtype>*>* bottom) {
   const Dtype* top_diff = top[0]->gpu_diff();
   Dtype* bottom_diff = (*bottom)[0]->mutable_gpu_diff();
-  caffe_gpu_copy(channels_out_, top_diff, bottom_diff);
+  caffe_gpu_copy(count_, top_diff, bottom_diff);
 }
 
 INSTANTIATE_CLASS(FlattenLayer);

diff --git a/src/caffe/test/test_flatten_layer.cpp b/src/caffe/test/test_flatten_layer.cpp
@@ -58,6 +58,8 @@ TYPED_TEST(FlattenLayerTest, TestCPU) {
   for (int c = 0; c < 3 * 6 * 5; ++c) {
     EXPECT_EQ(this->blob_top_->data_at(0, c, 0, 0),
         this->blob_bottom_->data_at(0, c / (6 * 5), (c / 5) % 6, c % 5));
+    EXPECT_EQ(this->blob_top_->data_at(1, c, 0, 0),
+        this->blob_bottom_->data_at(1, c / (6 * 5), (c / 5) % 6, c % 5));
   }
 }