replace std::shuffle with version using prefetch rng; improve unit test

diff --git a/include/caffe/vision_layers.hpp b/include/caffe/vision_layers.hpp
index 2296c8c..991e338 100644 (file)
--- a/include/caffe/vision_layers.hpp
+++ b/include/caffe/vision_layers.hpp
@@ -472,6 +472,8 @@ class ImageDataLayer : public Layer<Dtype> {
   virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
       const bool propagate_down, vector<Blob<Dtype>*>* bottom) { return; }
 
+  virtual void ShuffleImages();
+
   virtual void CreatePrefetchThread();
   virtual void JoinPrefetchThread();
   virtual unsigned int PrefetchRand();

diff --git a/src/caffe/layers/image_data_layer.cpp b/src/caffe/layers/image_data_layer.cpp
index bd56ef7..ed064d0 100644 (file)
--- a/src/caffe/layers/image_data_layer.cpp
+++ b/src/caffe/layers/image_data_layer.cpp
@@ -8,6 +8,7 @@
 #include <vector>
 #include <iostream>  // NOLINT(readability/streams)
 #include <fstream>  // NOLINT(readability/streams)
+#include <utility>
 
 #include "caffe/layer.hpp"
 #include "caffe/util/io.hpp"
@@ -15,6 +16,7 @@
 #include "caffe/util/rng.hpp"
 #include "caffe/vision_layers.hpp"
 
+using std::iterator;
 using std::string;
 using std::pair;
 
@@ -122,7 +124,7 @@ void* ImageDataLayerPrefetch(void* layer_pointer) {
       DLOG(INFO) << "Restarting data prefetching from start.";
       layer->lines_id_ = 0;
       if (layer->layer_param_.image_data_param().shuffle()) {
-        std::random_shuffle(layer->lines_.begin(), layer->lines_.end());
+        layer->ShuffleImages();
       }
     }
   }
@@ -158,7 +160,9 @@ void ImageDataLayer<Dtype>::SetUp(const vector<Blob<Dtype>*>& bottom,
   if (this->layer_param_.image_data_param().shuffle()) {
     // randomly shuffle data
     LOG(INFO) << "Shuffling data";
-    std::random_shuffle(lines_.begin(), lines_.end());
+    const unsigned int prefetch_rng_seed = caffe_rng_rand();
+    prefetch_rng_.reset(new Caffe::RNG(prefetch_rng_seed));
+    ShuffleImages();
   }
   LOG(INFO) << "A total of " << lines_.size() << " images.";
 
@@ -231,9 +235,11 @@ void ImageDataLayer<Dtype>::SetUp(const vector<Blob<Dtype>*>& bottom,
 template <typename Dtype>
 void ImageDataLayer<Dtype>::CreatePrefetchThread() {
   phase_ = Caffe::phase();
-  const bool prefetch_needs_rand = (phase_ == Caffe::TRAIN) &&
-      (this->layer_param_.data_param().mirror() ||
-       this->layer_param_.data_param().crop_size());
+  const bool prefetch_needs_rand =
+      this->layer_param_.image_data_param().shuffle() ||
+          ((phase_ == Caffe::TRAIN) &&
+           (this->layer_param_.image_data_param().mirror() ||
+            this->layer_param_.image_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));
@@ -246,6 +252,18 @@ void ImageDataLayer<Dtype>::CreatePrefetchThread() {
 }
 
 template <typename Dtype>
+void ImageDataLayer<Dtype>::ShuffleImages() {
+  const int num_images = lines_.size();
+  for (int i = 0; i < num_images; ++i) {
+    const int max_rand_index = num_images - i;
+    const int rand_index = PrefetchRand() % max_rand_index;
+    pair<string, int> item = lines_[rand_index];
+    lines_.erase(lines_.begin() + rand_index);
+    lines_.push_back(item);
+  }
+}
+
+template <typename Dtype>
 void ImageDataLayer<Dtype>::JoinPrefetchThread() {
   CHECK(!pthread_join(thread_, NULL)) << "Pthread joining failed.";
 }

diff --git a/src/caffe/layers/window_data_layer.cpp b/src/caffe/layers/window_data_layer.cpp
index 872bffb..862c034 100644 (file)
--- a/src/caffe/layers/window_data_layer.cpp
+++ b/src/caffe/layers/window_data_layer.cpp
@@ -410,8 +410,9 @@ void WindowDataLayer<Dtype>::SetUp(const vector<Blob<Dtype>*>& bottom,
 
 template <typename Dtype>
 void WindowDataLayer<Dtype>::CreatePrefetchThread() {
-  const bool prefetch_needs_rand = this->layer_param_.data_param().mirror() ||
-      this->layer_param_.data_param().crop_size();
+  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));

diff --git a/src/caffe/test/test_image_data_layer.cpp b/src/caffe/test/test_image_data_layer.cpp
index c300025..42a1d03 100644 (file)
--- a/src/caffe/test/test_image_data_layer.cpp
+++ b/src/caffe/test/test_image_data_layer.cpp
@@ -4,6 +4,7 @@
 
 #include <iostream>  // NOLINT(readability/streams)
 #include <fstream>  // NOLINT(readability/streams)
+#include <map>
 #include <string>
 #include <vector>
 
@@ -15,6 +16,7 @@
 #include "caffe/proto/caffe.pb.h"
 #include "caffe/test/test_caffe_main.hpp"
 
+using std::map;
 using std::string;
 
 namespace caffe {
@@ -27,14 +29,15 @@ class ImageDataLayerTest : public ::testing::Test {
   ImageDataLayerTest()
       : blob_top_data_(new Blob<Dtype>()),
         blob_top_label_(new Blob<Dtype>()),
-        filename(NULL) {}
+        filename_(new string(tmpnam(NULL))),
+        seed_(1701) {}
   virtual void SetUp() {
     blob_top_vec_.push_back(blob_top_data_);
     blob_top_vec_.push_back(blob_top_label_);
+    Caffe::set_random_seed(seed_);
     // Create a Vector of files with labels
-    filename = tmpnam(NULL);  // get temp name
-    std::ofstream outfile(filename, std::ofstream::out);
-    LOG(INFO) << "Using temporary file " << filename;
+    std::ofstream outfile(filename_->c_str(), std::ofstream::out);
+    LOG(INFO) << "Using temporary file " << *filename_;
     for (int i = 0; i < 5; ++i) {
       outfile << "examples/images/cat.jpg " << i;
     }
@@ -46,7 +49,8 @@ class ImageDataLayerTest : public ::testing::Test {
     delete blob_top_label_;
   }
 
-  char* filename;
+  int seed_;
+  shared_ptr<string> filename_;
   Blob<Dtype>* const blob_top_data_;
   Blob<Dtype>* const blob_top_label_;
   vector<Blob<Dtype>*> blob_bottom_vec_;
@@ -60,7 +64,7 @@ TYPED_TEST(ImageDataLayerTest, TestRead) {
   LayerParameter param;
   ImageDataParameter* image_data_param = param.mutable_image_data_param();
   image_data_param->set_batch_size(5);
-  image_data_param->set_source(this->filename);
+  image_data_param->set_source(this->filename_->c_str());
   image_data_param->set_shuffle(false);
   ImageDataLayer<TypeParam> layer(param);
   layer.SetUp(this->blob_bottom_vec_, &this->blob_top_vec_);
@@ -85,7 +89,7 @@ TYPED_TEST(ImageDataLayerTest, TestResize) {
   LayerParameter param;
   ImageDataParameter* image_data_param = param.mutable_image_data_param();
   image_data_param->set_batch_size(5);
-  image_data_param->set_source(this->filename);
+  image_data_param->set_source(this->filename_->c_str());
   image_data_param->set_new_height(256);
   image_data_param->set_new_width(256);
   image_data_param->set_shuffle(false);
@@ -112,7 +116,7 @@ TYPED_TEST(ImageDataLayerTest, TestShuffle) {
   LayerParameter param;
   ImageDataParameter* image_data_param = param.mutable_image_data_param();
   image_data_param->set_batch_size(5);
-  image_data_param->set_source(this->filename);
+  image_data_param->set_source(this->filename_->c_str());
   image_data_param->set_shuffle(true);
   ImageDataLayer<TypeParam> layer(param);
   layer.SetUp(this->blob_bottom_vec_, &this->blob_top_vec_);
@@ -127,10 +131,17 @@ TYPED_TEST(ImageDataLayerTest, TestShuffle) {
   // Go through the data twice
   for (int iter = 0; iter < 2; ++iter) {
     layer.Forward(this->blob_bottom_vec_, &this->blob_top_vec_);
+    map<TypeParam, int> values_to_indices;
+    int num_in_order = 0;
     for (int i = 0; i < 5; ++i) {
-      EXPECT_GE(this->blob_top_label_->cpu_data()[i], 0);
-      EXPECT_LE(this->blob_top_label_->cpu_data()[i], 5);
+      TypeParam value = this->blob_top_label_->cpu_data()[i];
+      // Check that the value has not been seen already (no duplicates).
+      EXPECT_EQ(values_to_indices.find(value), values_to_indices.end());
+      values_to_indices[value] = i;
+      num_in_order += (value == TypeParam(i));
     }
+    EXPECT_EQ(5, values_to_indices.size());
+    EXPECT_GT(5, num_in_order);
   }
 }