Reshape single input batches for inputs of varying dimension

src/caffe/layers/base_data_layer.cpp

@@ -61,6 +61,9 @@ void BasePrefetchingDataLayer<Dtype>::Forward_cpu(
   // First, join the thread
   JoinPrefetchThread();
   DLOG(INFO) << "Thread joined";
+  // Reshape to loaded data.
+  top[0]->Reshape(this->prefetch_data_.num(), this->prefetch_data_.channels(),
+      this->prefetch_data_.height(), this->prefetch_data_.width());
   // Copy the data
   caffe_copy(prefetch_data_.count(), prefetch_data_.cpu_data(),
              top[0]->mutable_cpu_data());

src/caffe/layers/base_data_layer.cu

@@ -9,6 +9,9 @@ void BasePrefetchingDataLayer<Dtype>::Forward_gpu(
     const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
   // First, join the thread
   JoinPrefetchThread();
+  // Reshape to loaded data.
+  top[0]->Reshape(this->prefetch_data_.num(), this->prefetch_data_.channels(),
+      this->prefetch_data_.height(), this->prefetch_data_.width());
   // Copy the data
   caffe_copy(prefetch_data_.count(), prefetch_data_.cpu_data(),
       top[0]->mutable_gpu_data());

src/caffe/layers/data_layer.cpp

@@ -49,7 +49,7 @@ void DataLayer<Dtype>::DataLayerSetUp(const vector<Blob<Dtype>*>& bottom,
   int crop_size = this->layer_param_.transform_param().crop_size();
   if (crop_size > 0) {
     top[0]->Reshape(this->layer_param_.data_param().batch_size(),
-                       datum.channels(), crop_size, crop_size);
+        datum.channels(), crop_size, crop_size);
     this->prefetch_data_.Reshape(this->layer_param_.data_param().batch_size(),
         datum.channels(), crop_size, crop_size);
     this->transformed_data_.Reshape(1, datum.channels(), crop_size, crop_size);
@@ -83,13 +83,25 @@ void DataLayer<Dtype>::InternalThreadEntry() {
   CPUTimer timer;
   CHECK(this->prefetch_data_.count());
   CHECK(this->transformed_data_.count());
+
+  // Reshape on single input batches for inputs of varying dimension.
+  const int batch_size = this->layer_param_.data_param().batch_size();
+  const int crop_size = this->layer_param_.transform_param().crop_size();
+  if (batch_size == 1 && crop_size == 0) {
+    Datum datum;
+    datum.ParseFromString(cursor_->value());
+    this->prefetch_data_.Reshape(1, datum.channels(),
+        datum.height(), datum.width());
+    this->transformed_data_.Reshape(1, datum.channels(),
+        datum.height(), datum.width());
+  }
+
   Dtype* top_data = this->prefetch_data_.mutable_cpu_data();
   Dtype* top_label = NULL;  // suppress warnings about uninitialized variables
 
   if (this->output_labels_) {
     top_label = this->prefetch_label_.mutable_cpu_data();
   }
-  const int batch_size = this->layer_param_.data_param().batch_size();
   for (int item_id = 0; item_id < batch_size; ++item_id) {
     timer.Start();
     // get a blob

src/caffe/layers/image_data_layer.cpp

@@ -102,15 +102,27 @@ void ImageDataLayer<Dtype>::InternalThreadEntry() {
   CPUTimer timer;
   CHECK(this->prefetch_data_.count());
   CHECK(this->transformed_data_.count());
-  Dtype* top_data = this->prefetch_data_.mutable_cpu_data();
-  Dtype* top_label = this->prefetch_label_.mutable_cpu_data();
   ImageDataParameter image_data_param = this->layer_param_.image_data_param();
   const int batch_size = image_data_param.batch_size();
   const int new_height = image_data_param.new_height();
   const int new_width = image_data_param.new_width();
+  const int crop_size = this->layer_param_.transform_param().crop_size();
   const bool is_color = image_data_param.is_color();
   string root_folder = image_data_param.root_folder();
 
+  // Reshape on single input batches for inputs of varying dimension.
+  if (batch_size == 1 && crop_size == 0 && new_height == 0 && new_width == 0) {
+    cv::Mat cv_img = ReadImageToCVMat(root_folder + lines_[lines_id_].first,
+        0, 0, is_color);
+    this->prefetch_data_.Reshape(1, cv_img.channels(),
+        cv_img.rows, cv_img.cols);
+    this->transformed_data_.Reshape(1, cv_img.channels(),
+        cv_img.rows, cv_img.cols);
+  }
+
+  Dtype* prefetch_data = this->prefetch_data_.mutable_cpu_data();
+  Dtype* prefetch_label = this->prefetch_label_.mutable_cpu_data();
+
   // datum scales
   const int lines_size = lines_.size();
   for (int item_id = 0; item_id < batch_size; ++item_id) {
@@ -124,11 +136,11 @@ void ImageDataLayer<Dtype>::InternalThreadEntry() {
     timer.Start();
     // Apply transformations (mirror, crop...) to the image
     int offset = this->prefetch_data_.offset(item_id);
-    this->transformed_data_.set_cpu_data(top_data + offset);
+    this->transformed_data_.set_cpu_data(prefetch_data + offset);
     this->data_transformer_.Transform(cv_img, &(this->transformed_data_));
     trans_time += timer.MicroSeconds();
 
-    top_label[item_id] = lines_[lines_id_].second;
+    prefetch_label[item_id] = lines_[lines_id_].second;
     // go to the next iter
     lines_id_++;
     if (lines_id_ >= lines_size) {

src/caffe/test/test_data_layer.cpp

@@ -103,6 +103,70 @@ class DataLayerTest : public MultiDeviceTest<TypeParam> {
     }
   }
 
+  void TestReshape(DataParameter_DB backend) {
+    const int num_inputs = 5;
+    // Save data of varying shapes.
+    LOG(INFO) << "Using temporary dataset " << *filename_;
+    scoped_ptr<db::DB> db(db::GetDB(backend));
+    db->Open(*filename_, db::NEW);
+    scoped_ptr<db::Transaction> txn(db->NewTransaction());
+    for (int i = 0; i < num_inputs; ++i) {
+      Datum datum;
+      datum.set_label(i);
+      datum.set_channels(2);
+      datum.set_height(i % 2 + 1);
+      datum.set_width(i % 4 + 1);
+      std::string* data = datum.mutable_data();
+      const int data_size = datum.channels() * datum.height() * datum.width();
+      for (int j = 0; j < data_size; ++j) {
+        data->push_back(static_cast<uint8_t>(j));
+      }
+      stringstream ss;
+      ss << i;
+      string out;
+      CHECK(datum.SerializeToString(&out));
+      txn->Put(ss.str(), out);
+    }
+    txn->Commit();
+    db->Close();
+
+    // Load and check data of various shapes.
+    LayerParameter param;
+    DataParameter* data_param = param.mutable_data_param();
+    data_param->set_batch_size(1);
+    data_param->set_source(filename_->c_str());
+    data_param->set_backend(backend);
+
+    DataLayer<Dtype> layer(param);
+    layer.SetUp(blob_bottom_vec_, blob_top_vec_);
+    EXPECT_EQ(blob_top_data_->num(), 1);
+    EXPECT_EQ(blob_top_data_->channels(), 2);
+    EXPECT_EQ(blob_top_label_->num(), 1);
+    EXPECT_EQ(blob_top_label_->channels(), 1);
+    EXPECT_EQ(blob_top_label_->height(), 1);
+    EXPECT_EQ(blob_top_label_->width(), 1);
+
+    for (int iter = 0; iter < num_inputs; ++iter) {
+      layer.Forward(blob_bottom_vec_, blob_top_vec_);
+      EXPECT_EQ(blob_top_data_->height(), iter % 2 + 1);
+      EXPECT_EQ(blob_top_data_->width(), iter % 4 + 1);
+      EXPECT_EQ(iter, blob_top_label_->cpu_data()[0]);
+      const int channels = blob_top_data_->channels();
+      const int height = blob_top_data_->height();
+      const int width = blob_top_data_->width();
+      for (int c = 0; c < channels; ++c) {
+        for (int h = 0; h < height; ++h) {
+          for (int w = 0; w < width; ++w) {
+            const int idx = (c * height + h) * width + w;
+            EXPECT_EQ(idx, static_cast<int>(blob_top_data_->cpu_data()[idx]))
+                << "debug: iter " << iter << " c " << c
+                << " h " << h << " w " << w;
+          }
+        }
+      }
+    }
+  }
+
   void TestReadCrop() {
     const Dtype scale = 3;
     LayerParameter param;
@@ -285,6 +349,10 @@ TYPED_TEST(DataLayerTest, TestReadLevelDB) {
   this->TestRead();
 }
 
+TYPED_TEST(DataLayerTest, TestReshapeLevelDB) {
+  this->TestReshape(DataParameter_DB_LEVELDB);
+}
+
 TYPED_TEST(DataLayerTest, TestReadCropTrainLevelDB) {
   Caffe::set_phase(Caffe::TRAIN);
   const bool unique_pixels = true;  // all images the same; pixels different
@@ -323,6 +391,10 @@ TYPED_TEST(DataLayerTest, TestReadLMDB) {
   this->TestRead();
 }
 
+TYPED_TEST(DataLayerTest, TestReshapeLMDB) {
+  this->TestReshape(DataParameter_DB_LMDB);
+}
+
 TYPED_TEST(DataLayerTest, TestReadCropTrainLMDB) {
   Caffe::set_phase(Caffe::TRAIN);
   const bool unique_pixels = true;  // all images the same; pixels different

src/caffe/test/test_image_data_layer.cpp

@@ -25,17 +25,24 @@ class ImageDataLayerTest : public MultiDeviceTest<TypeParam> {
         blob_top_data_(new Blob<Dtype>()),
         blob_top_label_(new Blob<Dtype>()) {}
   virtual void SetUp() {
-    MakeTempFilename(&filename_);
     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
+    // Create test input file.
+    MakeTempFilename(&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_SOURCE_DIR "images/cat.jpg " << i;
     }
     outfile.close();
+    // Create test input file for images of distinct sizes.
+    MakeTempFilename(&filename_reshape_);
+    std::ofstream reshapefile(filename_reshape_.c_str(), std::ofstream::out);
+    LOG(INFO) << "Using temporary file " << filename_reshape_;
+    reshapefile << EXAMPLES_SOURCE_DIR "images/cat.jpg " << 0;
+    reshapefile << EXAMPLES_SOURCE_DIR "images/fish-bike.jpg " << 1;
+    reshapefile.close();
   }
 
   virtual ~ImageDataLayerTest() {
@@ -45,6 +52,7 @@ class ImageDataLayerTest : public MultiDeviceTest<TypeParam> {
 
   int seed_;
   string filename_;
+  string filename_reshape_;
   Blob<Dtype>* const blob_top_data_;
   Blob<Dtype>* const blob_top_label_;
   vector<Blob<Dtype>*> blob_bottom_vec_;
@@ -107,6 +115,33 @@ TYPED_TEST(ImageDataLayerTest, TestResize) {
   }
 }
 
+TYPED_TEST(ImageDataLayerTest, TestReshape) {
+  typedef typename TypeParam::Dtype Dtype;
+  LayerParameter param;
+  ImageDataParameter* image_data_param = param.mutable_image_data_param();
+  image_data_param->set_batch_size(1);
+  image_data_param->set_source(this->filename_reshape_.c_str());
+  image_data_param->set_shuffle(false);
+  ImageDataLayer<Dtype> layer(param);
+  layer.SetUp(this->blob_bottom_vec_, this->blob_top_vec_);
+  EXPECT_EQ(this->blob_top_label_->num(), 1);
+  EXPECT_EQ(this->blob_top_label_->channels(), 1);
+  EXPECT_EQ(this->blob_top_label_->height(), 1);
+  EXPECT_EQ(this->blob_top_label_->width(), 1);
+  // cat.jpg
+  layer.Forward(this->blob_bottom_vec_, this->blob_top_vec_);
+  EXPECT_EQ(this->blob_top_data_->num(), 1);
+  EXPECT_EQ(this->blob_top_data_->channels(), 3);
+  EXPECT_EQ(this->blob_top_data_->height(), 360);
+  EXPECT_EQ(this->blob_top_data_->width(), 480);
+  // fish-bike.jpg
+  layer.Forward(this->blob_bottom_vec_, this->blob_top_vec_);
+  EXPECT_EQ(this->blob_top_data_->num(), 1);
+  EXPECT_EQ(this->blob_top_data_->channels(), 3);
+  EXPECT_EQ(this->blob_top_data_->height(), 323);
+  EXPECT_EQ(this->blob_top_data_->width(), 481);
+}
+
 TYPED_TEST(ImageDataLayerTest, TestShuffle) {
   typedef typename TypeParam::Dtype Dtype;
   LayerParameter param;