make HDF5 layer support multiple data output

diff --git a/include/caffe/data_layers.hpp b/include/caffe/data_layers.hpp
index b3f93ce..c1720b7 100644 (file)
--- a/include/caffe/data_layers.hpp
+++ b/include/caffe/data_layers.hpp
@@ -20,9 +20,6 @@
 
 namespace caffe {
 
-#define HDF5_DATA_DATASET_NAME "data"
-#define HDF5_DATA_LABEL_NAME "label"
-
 /**
  * @brief Provides base for data layers that feed blobs to the Net.
  *
@@ -180,7 +177,7 @@ class HDF5DataLayer : public Layer<Dtype> {
     return LayerParameter_LayerType_HDF5_DATA;
   }
   virtual inline int ExactNumBottomBlobs() const { return 0; }
-  virtual inline int ExactNumTopBlobs() const { return 2; }
+  virtual inline int MinTopBlobs() const { return 1; }
 
  protected:
   virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
@@ -197,8 +194,7 @@ class HDF5DataLayer : public Layer<Dtype> {
   unsigned int num_files_;
   unsigned int current_file_;
   hsize_t current_row_;
-  Blob<Dtype> data_blob_;
-  Blob<Dtype> label_blob_;
+  std::vector<shared_ptr<Blob<Dtype> > > hdf_blobs_;
 };
 
 /**

diff --git a/src/caffe/layers/hdf5_data_layer.cpp b/src/caffe/layers/hdf5_data_layer.cpp
index 8fec617..bf771fd 100644 (file)
--- a/src/caffe/layers/hdf5_data_layer.cpp
+++ b/src/caffe/layers/hdf5_data_layer.cpp
@@ -33,20 +33,27 @@ void HDF5DataLayer<Dtype>::LoadHDF5FileData(const char* filename) {
     return;
   }
 
-  const int MIN_DATA_DIM = 2;
+  int top_size = this->layer_param_.top_size();
+  hdf_blobs_.resize(top_size);
+
+  const int MIN_DATA_DIM = 1;
   const int MAX_DATA_DIM = 4;
-  hdf5_load_nd_dataset(
-    file_id, "data",  MIN_DATA_DIM, MAX_DATA_DIM, &data_blob_);
 
-  const int MIN_LABEL_DIM = 1;
-  const int MAX_LABEL_DIM = 2;
-  hdf5_load_nd_dataset(
-    file_id, "label", MIN_LABEL_DIM, MAX_LABEL_DIM, &label_blob_);
+  for (int i = 0; i < top_size; ++i) {
+      hdf_blobs_[i] = shared_ptr<Blob<Dtype> >(new Blob<Dtype>());
+      hdf5_load_nd_dataset(file_id, this->layer_param_.top(i).c_str(),
+        MIN_DATA_DIM, MAX_DATA_DIM, hdf_blobs_[i].get());
+  }
 
   herr_t status = H5Fclose(file_id);
   CHECK_GE(status, 0) << "Failed to close HDF5 file " << filename;
-  CHECK_EQ(data_blob_.num(), label_blob_.num());
-  LOG(INFO) << "Successully loaded " << data_blob_.num() << " rows";
+
+  // MinTopBlobs==1 guarantees at least one top blob
+  int num = hdf_blobs_[0]->num();
+  for (int i = 1; i < top_size; ++i) {
+      CHECK_EQ(hdf_blobs_[i]->num(), num);
+  }
+  LOG(INFO) << "Successully loaded " << hdf_blobs_[0]->num() << " rows";
 }
 
 template <typename Dtype>
@@ -74,26 +81,21 @@ void HDF5DataLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
 
   // Reshape blobs.
   const int batch_size = this->layer_param_.hdf5_data_param().batch_size();
-  top[0]->Reshape(batch_size, data_blob_.channels(),
-                  data_blob_.height(), data_blob_.width());
-  top[1]->Reshape(batch_size, label_blob_.channels(),
-                  label_blob_.height(), label_blob_.width());
-  LOG(INFO) << "output data size: " << top[0]->num() << ","
-      << top[0]->channels() << "," << top[0]->height() << ","
-      << top[0]->width();
+  const int top_size = this->layer_param_.top_size();
+  for (int i = 0; i < top_size; ++i) {
+    top[i]->Reshape(batch_size, hdf_blobs_[i]->channels(),
+                    hdf_blobs_[i]->height(), hdf_blobs_[i]->width());
+  }
 }
 
 template <typename Dtype>
 void HDF5DataLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
       const vector<Blob<Dtype>*>& top) {
   const int batch_size = this->layer_param_.hdf5_data_param().batch_size();
-  const int data_count = top[0]->count() / top[0]->num();
-  const int label_data_count = top[1]->count() / top[1]->num();
-
   for (int i = 0; i < batch_size; ++i, ++current_row_) {
-    if (current_row_ == data_blob_.num()) {
+    if (current_row_ == hdf_blobs_[0]->num()) {
       if (num_files_ > 1) {
-        current_file_ += 1;
+        ++current_file_;
         if (current_file_ == num_files_) {
           current_file_ = 0;
           LOG(INFO) << "looping around to first file";
@@ -102,11 +104,12 @@ void HDF5DataLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
       }
       current_row_ = 0;
     }
-    caffe_copy(data_count, &data_blob_.cpu_data()[current_row_ * data_count],
-               &top[0]->mutable_cpu_data()[i * data_count]);
-    caffe_copy(label_data_count,
-               &label_blob_.cpu_data()[current_row_ * label_data_count],
-               &top[1]->mutable_cpu_data()[i * label_data_count]);
+    for (int j = 0; j < this->layer_param_.top_size(); ++j) {
+      int data_dim = top[j]->count() / top[j]->num();
+      caffe_copy(data_dim,
+          &hdf_blobs_[j]->cpu_data()[current_row_ * data_dim],
+          &top[j]->mutable_cpu_data()[i * data_dim]);
+    }
   }
 }
 

diff --git a/src/caffe/layers/hdf5_data_layer.cu b/src/caffe/layers/hdf5_data_layer.cu
index f671ea1..02ec673 100644 (file)
--- a/src/caffe/layers/hdf5_data_layer.cu
+++ b/src/caffe/layers/hdf5_data_layer.cu
@@ -20,29 +20,24 @@ template <typename Dtype>
 void HDF5DataLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype>*>& bottom,
       const vector<Blob<Dtype>*>& top) {
   const int batch_size = this->layer_param_.hdf5_data_param().batch_size();
-  const int data_count = top[0]->count() / top[0]->num();
-  const int label_data_count = top[1]->count() / top[1]->num();
-
   for (int i = 0; i < batch_size; ++i, ++current_row_) {
-    if (current_row_ == data_blob_.num()) {
+    if (current_row_ == hdf_blobs_[0]->num()) {
       if (num_files_ > 1) {
         current_file_ += 1;
-
         if (current_file_ == num_files_) {
           current_file_ = 0;
           LOG(INFO) << "looping around to first file";
         }
-
         LoadHDF5FileData(hdf_filenames_[current_file_].c_str());
       }
       current_row_ = 0;
     }
-    caffe_copy(data_count,
-        &data_blob_.cpu_data()[current_row_ * data_count],
-        &top[0]->mutable_gpu_data()[i * data_count]);
-    caffe_copy(label_data_count,
-        &label_blob_.cpu_data()[current_row_ * label_data_count],
-        &top[1]->mutable_gpu_data()[i * label_data_count]);
+    for (int j = 0; j < this->layer_param_.top_size(); ++j) {
+      int data_dim = top[j]->count() / top[j]->num();
+      caffe_copy(data_dim,
+          &hdf_blobs_[j]->cpu_data()[current_row_ * data_dim],
+          &top[j]->mutable_gpu_data()[i * data_dim]);
+    }
   }
 }
 

diff --git a/src/caffe/test/test_hdf5data_layer.cpp b/src/caffe/test/test_hdf5data_layer.cpp
index 9026a91..769c132 100644 (file)
--- a/src/caffe/test/test_hdf5data_layer.cpp
+++ b/src/caffe/test/test_hdf5data_layer.cpp
@@ -57,6 +57,9 @@ TYPED_TEST(HDF5DataLayerTest, TestRead) {
   // The data file we are reading has 10 rows and 8 columns,
   // with values from 0 to 10*8 reshaped in row-major order.
   LayerParameter param;
+  param.add_top("data");
+  param.add_top("label");
+
   HDF5DataParameter* hdf5_data_param = param.mutable_hdf5_data_param();
   int batch_size = 5;
   hdf5_data_param->set_batch_size(batch_size);