Simplify the WindowDataLayer using the base class

diff --git a/include/caffe/common.hpp b/include/caffe/common.hpp
index 9008eae..683d1d6 100644 (file)
--- a/include/caffe/common.hpp
+++ b/include/caffe/common.hpp
@@ -6,9 +6,14 @@
 #include <glog/logging.h>
 
 #include <cmath>
+#include <fstream>  // NOLINT(readability/streams)
+#include <iostream>  // NOLINT(readability/streams)
 #include <map>
 #include <set>
+#include <sstream>
 #include <string>
+#include <utility>  // pair
+#include <vector>
 
 #include "caffe/util/device_alternate.hpp"
 
@@ -88,6 +93,7 @@ using std::ostringstream;
 using std::pair;
 using std::set;
 using std::string;
+using std::stringstream;
 using std::vector;
 
 // A global initialization function that you should call in your main function.

diff --git a/include/caffe/data_layers.hpp b/include/caffe/data_layers.hpp
index 499ee7f..8651ece 100644 (file)
--- a/include/caffe/data_layers.hpp
+++ b/include/caffe/data_layers.hpp
@@ -269,12 +269,12 @@ class MemoryDataLayer : public Layer<Dtype> {
 };
 
 template <typename Dtype>
-class WindowDataLayer : public Layer<Dtype>, public InternalThread {
+class WindowDataLayer : public BasePrefetchingDataLayer<Dtype> {
  public:
   explicit WindowDataLayer(const LayerParameter& param)
-      : Layer<Dtype>(param) {}
-  virtual ~WindowDataLayer();
-  virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
+      : BasePrefetchingDataLayer<Dtype>(param) {}
+  virtual ~WindowDataLayer() {}
+  virtual void DataLayerSetUp(const vector<Blob<Dtype>*>& bottom,
       vector<Blob<Dtype>*>* top);
 
   virtual inline LayerParameter_LayerType type() const {
@@ -284,24 +284,10 @@ class WindowDataLayer : public Layer<Dtype>, public InternalThread {
   virtual inline int ExactNumTopBlobs() const { return 2; }
 
  protected:
-  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
-      vector<Blob<Dtype>*>* top);
-  virtual void Forward_gpu(const vector<Blob<Dtype>*>& bottom,
-      vector<Blob<Dtype>*>* top);
-  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
-      const vector<bool>& propagate_down, vector<Blob<Dtype>*>* bottom) {}
-  virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
-      const vector<bool>& propagate_down, vector<Blob<Dtype>*>* bottom) {}
-
-  virtual void CreatePrefetchThread();
-  virtual void JoinPrefetchThread();
   virtual unsigned int PrefetchRand();
   virtual void InternalThreadEntry();
 
   shared_ptr<Caffe::RNG> prefetch_rng_;
-  Blob<Dtype> prefetch_data_;
-  Blob<Dtype> prefetch_label_;
-  Blob<Dtype> data_mean_;
   vector<std::pair<std::string, vector<int> > > image_database_;
   enum WindowField { IMAGE_INDEX, LABEL, OVERLAP, X1, Y1, X2, Y2, NUM };
   vector<vector<float> > fg_windows_;

diff --git a/src/caffe/layers/base_data_layer.cpp b/src/caffe/layers/base_data_layer.cpp
index 957a771..8310710 100644 (file)
--- a/src/caffe/layers/base_data_layer.cpp
+++ b/src/caffe/layers/base_data_layer.cpp
@@ -1,3 +1,5 @@
+#include <vector>
+
 #include "caffe/data_layers.hpp"
 
 namespace caffe {

diff --git a/src/caffe/layers/data_layer.cpp b/src/caffe/layers/data_layer.cpp
index 1b3acbd..fbf2d91 100644 (file)
--- a/src/caffe/layers/data_layer.cpp
+++ b/src/caffe/layers/data_layer.cpp
@@ -4,12 +4,13 @@
 #include <string>
 #include <vector>
 
+#include "caffe/common.hpp"
+#include "caffe/data_layers.hpp"
 #include "caffe/layer.hpp"
 #include "caffe/proto/caffe.pb.h"
 #include "caffe/util/io.hpp"
 #include "caffe/util/math_functions.hpp"
 #include "caffe/util/rng.hpp"
-#include "caffe/vision_layers.hpp"
 
 namespace caffe {
 

diff --git a/src/caffe/layers/image_data_layer.cpp b/src/caffe/layers/image_data_layer.cpp
index e4805ca..687beb8 100644 (file)
--- a/src/caffe/layers/image_data_layer.cpp
+++ b/src/caffe/layers/image_data_layer.cpp
@@ -4,11 +4,11 @@
 #include <utility>
 #include <vector>
 
+#include "caffe/data_layers.hpp"
 #include "caffe/layer.hpp"
 #include "caffe/util/io.hpp"
 #include "caffe/util/math_functions.hpp"
 #include "caffe/util/rng.hpp"
-#include "caffe/vision_layers.hpp"
 
 namespace caffe {
 
@@ -101,7 +101,8 @@ void ImageDataLayer<Dtype>::DataLayerSetUp(const vector<Blob<Dtype>*>& bottom,
       .transform_param().mean_file();
   if (crop_size > 0) {
     (*top)[0]->Reshape(batch_size, datum.channels(), crop_size, crop_size);
-    this->prefetch_data_.Reshape(batch_size, datum.channels(), crop_size, crop_size);
+    this->prefetch_data_.Reshape(batch_size, datum.channels(), crop_size,
+                                 crop_size);
   } else {
     (*top)[0]->Reshape(batch_size, datum.channels(), datum.height(),
                        datum.width());

diff --git a/src/caffe/layers/window_data_layer.cpp b/src/caffe/layers/window_data_layer.cpp
index ab12c61..f434982 100644 (file)
--- a/src/caffe/layers/window_data_layer.cpp
+++ b/src/caffe/layers/window_data_layer.cpp
@@ -1,7 +1,6 @@
 #include <stdint.h>
 
 #include <algorithm>
-#include <fstream>  // NOLINT(readability/streams)
 #include <map>
 #include <string>
 #include <utility>
@@ -11,11 +10,12 @@
 #include "opencv2/highgui/highgui.hpp"
 #include "opencv2/imgproc/imgproc.hpp"
 
+#include "caffe/common.hpp"
+#include "caffe/data_layers.hpp"
 #include "caffe/layer.hpp"
 #include "caffe/util/io.hpp"
 #include "caffe/util/math_functions.hpp"
 #include "caffe/util/rng.hpp"
-#include "caffe/vision_layers.hpp"
 
 // caffe.proto > LayerParameter > WindowDataParameter
 //   'source' field specifies the window_file
@@ -29,8 +29,8 @@ void WindowDataLayer<Dtype>::InternalThreadEntry() {
   // At each iteration, sample N windows where N*p are foreground (object)
   // windows and N*(1-p) are background (non-object) windows
 
-  Dtype* top_data = prefetch_data_.mutable_cpu_data();
-  Dtype* top_label = prefetch_label_.mutable_cpu_data();
+  Dtype* top_data = this->prefetch_data_.mutable_cpu_data();
+  Dtype* top_label = this->prefetch_label_.mutable_cpu_data();
   const Dtype scale = this->layer_param_.window_data_param().scale();
   const int batch_size = this->layer_param_.window_data_param().batch_size();
   const int crop_size = this->layer_param_.window_data_param().crop_size();
@@ -38,17 +38,17 @@ void WindowDataLayer<Dtype>::InternalThreadEntry() {
   const bool mirror = this->layer_param_.window_data_param().mirror();
   const float fg_fraction =
       this->layer_param_.window_data_param().fg_fraction();
-  const Dtype* mean = data_mean_.cpu_data();
-  const int mean_off = (data_mean_.width() - crop_size) / 2;
-  const int mean_width = data_mean_.width();
-  const int mean_height = data_mean_.height();
+  const Dtype* mean = this->data_mean_.cpu_data();
+  const int mean_off = (this->data_mean_.width() - crop_size) / 2;
+  const int mean_width = this->data_mean_.width();
+  const int mean_height = this->data_mean_.height();
   cv::Size cv_crop_size(crop_size, crop_size);
   const string& crop_mode = this->layer_param_.window_data_param().crop_mode();
 
   bool use_square = (crop_mode == "square") ? true : false;
 
   // zero out batch
-  caffe_set(prefetch_data_.count(), Dtype(0), top_data);
+  caffe_set(this->prefetch_data_.count(), Dtype(0), top_data);
 
   const int num_fg = static_cast<int>(static_cast<float>(batch_size)
       * fg_fraction);
@@ -238,12 +238,7 @@ void WindowDataLayer<Dtype>::InternalThreadEntry() {
 }
 
 template <typename Dtype>
-WindowDataLayer<Dtype>::~WindowDataLayer<Dtype>() {
-  JoinPrefetchThread();
-}
-
-template <typename Dtype>
-void WindowDataLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
+void WindowDataLayer<Dtype>::DataLayerSetUp(const vector<Blob<Dtype>*>& bottom,
       vector<Blob<Dtype>*>* top) {
   // LayerSetUp runs through the window_file and creates two structures
   // that hold windows: one for foreground (object) windows and one
@@ -268,6 +263,16 @@ void WindowDataLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
       << "  foreground sampling fraction: "
       << this->layer_param_.window_data_param().fg_fraction();
 
+  const bool prefetch_needs_rand =
+      this->layer_param_.window_data_param().mirror() ||
+      this->layer_param_.window_data_param().crop_size();
+  if (prefetch_needs_rand) {
+    const unsigned int prefetch_rng_seed = caffe_rng_rand();
+    prefetch_rng_.reset(new Caffe::RNG(prefetch_rng_seed));
+  } else {
+    prefetch_rng_.reset();
+  }
+
   std::ifstream infile(this->layer_param_.window_data_param().source().c_str());
   CHECK(infile.good()) << "Failed to open window file "
       << this->layer_param_.window_data_param().source() << std::endl;
@@ -357,14 +362,14 @@ void WindowDataLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
   CHECK_GT(crop_size, 0);
   const int batch_size = this->layer_param_.window_data_param().batch_size();
   (*top)[0]->Reshape(batch_size, channels, crop_size, crop_size);
-  prefetch_data_.Reshape(batch_size, channels, crop_size, crop_size);
+  this->prefetch_data_.Reshape(batch_size, channels, crop_size, crop_size);
 
   LOG(INFO) << "output data size: " << (*top)[0]->num() << ","
       << (*top)[0]->channels() << "," << (*top)[0]->height() << ","
       << (*top)[0]->width();
   // label
   (*top)[1]->Reshape(batch_size, 1, 1, 1);
-  prefetch_label_.Reshape(batch_size, 1, 1, 1);
+  this->prefetch_label_.Reshape(batch_size, 1, 1, 1);
 
   // check if we want to have mean
   if (this->layer_param_.window_data_param().has_mean_file()) {
@@ -373,47 +378,27 @@ void WindowDataLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
     LOG(INFO) << "Loading mean file from" << mean_file;
     BlobProto blob_proto;
     ReadProtoFromBinaryFileOrDie(mean_file, &blob_proto);
-    data_mean_.FromProto(blob_proto);
-    CHECK_EQ(data_mean_.num(), 1);
-    CHECK_EQ(data_mean_.width(), data_mean_.height());
-    CHECK_EQ(data_mean_.channels(), channels);
+    this->data_mean_.FromProto(blob_proto);
+    CHECK_EQ(this->data_mean_.num(), 1);
+    CHECK_EQ(this->data_mean_.width(), this->data_mean_.height());
+    CHECK_EQ(this->data_mean_.channels(), channels);
   } else {
     // Simply initialize an all-empty mean.
-    data_mean_.Reshape(1, channels, crop_size, crop_size);
+    this->data_mean_.Reshape(1, channels, crop_size, crop_size);
   }
   // Now, start the prefetch thread. Before calling prefetch, we make two
   // cpu_data calls so that the prefetch thread does not accidentally make
   // simultaneous cudaMalloc calls when the main thread is running. In some
   // GPUs this seems to cause failures if we do not so.
-  prefetch_data_.mutable_cpu_data();
-  prefetch_label_.mutable_cpu_data();
-  data_mean_.cpu_data();
+  this->prefetch_data_.mutable_cpu_data();
+  this->prefetch_label_.mutable_cpu_data();
+  this->data_mean_.cpu_data();
   DLOG(INFO) << "Initializing prefetch";
-  CreatePrefetchThread();
+  this->CreatePrefetchThread();
   DLOG(INFO) << "Prefetch initialized.";
 }
 
 template <typename Dtype>
-void WindowDataLayer<Dtype>::CreatePrefetchThread() {
-  const bool prefetch_needs_rand =
-      this->layer_param_.window_data_param().mirror() ||
-      this->layer_param_.window_data_param().crop_size();
-  if (prefetch_needs_rand) {
-    const unsigned int prefetch_rng_seed = caffe_rng_rand();
-    prefetch_rng_.reset(new Caffe::RNG(prefetch_rng_seed));
-  } else {
-    prefetch_rng_.reset();
-  }
-  // Create the thread.
-  CHECK(StartInternalThread()) << "Thread execution failed.";
-}
-
-template <typename Dtype>
-void WindowDataLayer<Dtype>::JoinPrefetchThread() {
-  CHECK(WaitForInternalThreadToExit()) << "Thread joining failed.";
-}
-
-template <typename Dtype>
 unsigned int WindowDataLayer<Dtype>::PrefetchRand() {
   CHECK(prefetch_rng_);
   caffe::rng_t* prefetch_rng =
@@ -421,24 +406,6 @@ unsigned int WindowDataLayer<Dtype>::PrefetchRand() {
   return (*prefetch_rng)();
 }
 
-template <typename Dtype>
-void WindowDataLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
-      vector<Blob<Dtype>*>* top) {
-  // First, join the thread
-  JoinPrefetchThread();
-  // Copy the data
-  caffe_copy(prefetch_data_.count(), prefetch_data_.cpu_data(),
-             (*top)[0]->mutable_cpu_data());
-  caffe_copy(prefetch_label_.count(), prefetch_label_.cpu_data(),
-             (*top)[1]->mutable_cpu_data());
-  // Start a new prefetch thread
-  CreatePrefetchThread();
-}
-
-#ifdef CPU_ONLY
-STUB_GPU_FORWARD(WindowDataLayer, Forward);
-#endif
-
 INSTANTIATE_CLASS(WindowDataLayer);
 
 }  // namespace caffe

diff --git a/src/caffe/test/test_data_layer.cpp b/src/caffe/test/test_data_layer.cpp
index 5c21f2e..52cc2c0 100644 (file)
--- a/src/caffe/test/test_data_layer.cpp
+++ b/src/caffe/test/test_data_layer.cpp
@@ -1,9 +1,8 @@
 #include <string>
 #include <vector>
 
-#include "leveldb/db.h"
-
 #include "gtest/gtest.h"
+#include "leveldb/db.h"
 
 #include "caffe/blob.hpp"
 #include "caffe/common.hpp"
@@ -13,9 +12,6 @@
 
 #include "caffe/test/test_caffe_main.hpp"
 
-using std::string;
-using std::stringstream;
-
 namespace caffe {
 
 template <typename TypeParam>

diff --git a/src/caffe/test/test_image_data_layer.cpp b/src/caffe/test/test_image_data_layer.cpp
index 73be5da..ab2365c 100644 (file)
--- a/src/caffe/test/test_image_data_layer.cpp
+++ b/src/caffe/test/test_image_data_layer.cpp
@@ -1,5 +1,3 @@
-#include <fstream>  // NOLINT(readability/streams)
-#include <iostream>  // NOLINT(readability/streams)
 #include <map>
 #include <string>
 #include <vector>
@@ -14,9 +12,6 @@
 
 #include "caffe/test/test_caffe_main.hpp"
 
-using std::map;
-using std::string;
-
 namespace caffe {
 
 template <typename TypeParam>

diff --git a/src/caffe/test/test_window_data_layer.cpp b/src/caffe/test/test_window_data_layer.cpp
new file mode 100644 (file)
index 0000000..47591b6
--- /dev/null
+++ b/src/caffe/test/test_window_data_layer.cpp
@@ -0,0 +1,49 @@
+#include <string>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "caffe/blob.hpp"
+#include "caffe/common.hpp"
+#include "caffe/data_layers.hpp"
+#include "caffe/proto/caffe.pb.h"
+
+#include "caffe/test/test_caffe_main.hpp"
+
+namespace caffe {
+
+template <typename TypeParam>
+class WindowDataLayerTest : public MultiDeviceTest<TypeParam> {
+  typedef typename TypeParam::Dtype Dtype;
+
+ protected:
+  WindowDataLayerTest()
+      : seed_(1701),
+        filename_(new string(tmpnam(NULL))),
+        blob_top_data_(new Blob<Dtype>()),
+        blob_top_label_(new Blob<Dtype>()) {}
+  virtual void SetUp() {
+    blob_top_vec_.push_back(blob_top_data_);
+    blob_top_vec_.push_back(blob_top_label_);
+    Caffe::set_random_seed(seed_);
+  }
+
+  virtual ~WindowDataLayerTest() {
+    delete blob_top_data_;
+    delete blob_top_label_;
+  }
+
+  int seed_;
+  shared_ptr<string> filename_;
+  Blob<Dtype>* const blob_top_data_;
+  Blob<Dtype>* const blob_top_label_;
+  vector<Blob<Dtype>*> blob_bottom_vec_;
+  vector<Blob<Dtype>*> blob_top_vec_;
+};
+
+TYPED_TEST_CASE(WindowDataLayerTest, TestDtypesAndDevices);
+
+TYPED_TEST(WindowDataLayerTest, TestNothing) {
+}
+
+}  // namespace caffe