src/caffe/layers/hdf5_data_layer.cpp

   1 #ifdef USE_HDF5
   2 /*
   3 TODO:
   4 - load file in a separate thread ("prefetch")
   5 - can be smarter about the memcpy call instead of doing it row-by-row
   6   :: use util functions caffe_copy, and Blob->offset()
   7   :: don't forget to update hdf5_daa_layer.cu accordingly
   8 - add ability to shuffle filenames if flag is set
   9 */
  10 #include <fstream>  // NOLINT(readability/streams)
  11 #include <string>
  12 #include <vector>
  13
  14 #include "hdf5.h"
  15 #include "hdf5_hl.h"
  16 #include "stdint.h"
  17
  18 #include "caffe/layers/hdf5_data_layer.hpp"
  19 #include "caffe/util/hdf5.hpp"
  20
  21 namespace caffe {
  22
  23 template <typename Dtype>
  24 HDF5DataLayer<Dtype>::~HDF5DataLayer<Dtype>() { }
  25
  26 // Load data and label from HDF5 filename into the class property blobs.
  27 template <typename Dtype>
  28 void HDF5DataLayer<Dtype>::LoadHDF5FileData(const char* filename) {
  29   DLOG(INFO) << "Loading HDF5 file: " << filename;
  30   hid_t file_id = H5Fopen(filename, H5F_ACC_RDONLY, H5P_DEFAULT);
  31   if (file_id < 0) {
  32     LOG(FATAL) << "Failed opening HDF5 file: " << filename;
  33   }
  34
  35   int top_size = this->layer_param_.top_size();
  36   hdf_blobs_.resize(top_size);
  37
  38   const int MIN_DATA_DIM = 1;
  39   const int MAX_DATA_DIM = INT_MAX;
  40
  41   for (int i = 0; i < top_size; ++i) {
  42     hdf_blobs_[i] = shared_ptr<Blob<Dtype> >(new Blob<Dtype>());
  43     // Allow reshape here, as we are loading data not params
  44     hdf5_load_nd_dataset(file_id, this->layer_param_.top(i).c_str(),
  45         MIN_DATA_DIM, MAX_DATA_DIM, hdf_blobs_[i].get(), true);
  46   }
  47
  48   herr_t status = H5Fclose(file_id);
  49   CHECK_GE(status, 0) << "Failed to close HDF5 file: " << filename;
  50
  51   // MinTopBlobs==1 guarantees at least one top blob
  52   CHECK_GE(hdf_blobs_[0]->num_axes(), 1) << "Input must have at least 1 axis.";
  53   const int num = hdf_blobs_[0]->shape(0);
  54   for (int i = 1; i < top_size; ++i) {
  55     CHECK_EQ(hdf_blobs_[i]->shape(0), num);
  56   }
  57   // Default to identity permutation.
  58   data_permutation_.clear();
  59   data_permutation_.resize(hdf_blobs_[0]->shape(0));
  60   for (int i = 0; i < hdf_blobs_[0]->shape(0); i++)
  61     data_permutation_[i] = i;
  62
  63   // Shuffle if needed.
  64   if (this->layer_param_.hdf5_data_param().shuffle()) {
  65     std::random_shuffle(data_permutation_.begin(), data_permutation_.end());
  66     DLOG(INFO) << "Successfully loaded " << hdf_blobs_[0]->shape(0)
  67                << " rows (shuffled)";
  68   } else {
  69     DLOG(INFO) << "Successfully loaded " << hdf_blobs_[0]->shape(0) << " rows";
  70   }
  71 }
  72
  73 template <typename Dtype>
  74 void HDF5DataLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
  75       const vector<Blob<Dtype>*>& top) {
  76   // Refuse transformation parameters since HDF5 is totally generic.
  77   CHECK(!this->layer_param_.has_transform_param()) <<
  78       this->type() << " does not transform data.";
  79   // Read the source to parse the filenames.
  80   const string& source = this->layer_param_.hdf5_data_param().source();
  81   LOG(INFO) << "Loading list of HDF5 filenames from: " << source;
  82   hdf_filenames_.clear();
  83   std::ifstream source_file(source.c_str());
  84   if (source_file.is_open()) {
  85     std::string line;
  86     while (source_file >> line) {
  87       hdf_filenames_.push_back(line);
  88     }
  89   } else {
  90     LOG(FATAL) << "Failed to open source file: " << source;
  91   }
  92   source_file.close();
  93   num_files_ = hdf_filenames_.size();
  94   current_file_ = 0;
  95   LOG(INFO) << "Number of HDF5 files: " << num_files_;
  96   CHECK_GE(num_files_, 1) << "Must have at least 1 HDF5 filename listed in "
  97     << source;
  98
  99   file_permutation_.clear();
 100   file_permutation_.resize(num_files_);
 101   // Default to identity permutation.
 102   for (int i = 0; i < num_files_; i++) {
 103     file_permutation_[i] = i;
 104   }
 105
 106   // Shuffle if needed.
 107   if (this->layer_param_.hdf5_data_param().shuffle()) {
 108     std::random_shuffle(file_permutation_.begin(), file_permutation_.end());
 109   }
 110
 111   // Load the first HDF5 file and initialize the line counter.
 112   LoadHDF5FileData(hdf_filenames_[file_permutation_[current_file_]].c_str());
 113   current_row_ = 0;
 114
 115   // Reshape blobs.
 116   const int batch_size = this->layer_param_.hdf5_data_param().batch_size();
 117   const int top_size = this->layer_param_.top_size();
 118   vector<int> top_shape;
 119   for (int i = 0; i < top_size; ++i) {
 120     top_shape.resize(hdf_blobs_[i]->num_axes());
 121     top_shape[0] = batch_size;
 122     for (int j = 1; j < top_shape.size(); ++j) {
 123       top_shape[j] = hdf_blobs_[i]->shape(j);
 124     }
 125     top[i]->Reshape(top_shape);
 126   }
 127 }
 128
 129 template <typename Dtype>
 130 bool HDF5DataLayer<Dtype>::Skip() {
 131   int size = Caffe::solver_count();
 132   int rank = Caffe::solver_rank();
 133   bool keep = (offset_ % size) == rank ||
 134               // In test mode, only rank 0 runs, so avoid skipping
 135               this->layer_param_.phase() == TEST;
 136   return !keep;
 137 }
 138
 139 template<typename Dtype>
 140 void HDF5DataLayer<Dtype>::Next() {
 141   if (++current_row_ == hdf_blobs_[0]->shape(0)) {
 142     if (num_files_ > 1) {
 143       ++current_file_;
 144       if (current_file_ == num_files_) {
 145         current_file_ = 0;
 146         if (this->layer_param_.hdf5_data_param().shuffle()) {
 147           std::random_shuffle(file_permutation_.begin(),
 148                               file_permutation_.end());
 149         }
 150         DLOG(INFO) << "Looping around to first file.";
 151       }
 152       LoadHDF5FileData(
 153         hdf_filenames_[file_permutation_[current_file_]].c_str());
 154     }
 155     current_row_ = 0;
 156     if (this->layer_param_.hdf5_data_param().shuffle())
 157       std::random_shuffle(data_permutation_.begin(), data_permutation_.end());
 158   }
 159   offset_++;
 160 }
 161
 162 template <typename Dtype>
 163 void HDF5DataLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
 164       const vector<Blob<Dtype>*>& top) {
 165   const int batch_size = this->layer_param_.hdf5_data_param().batch_size();
 166   for (int i = 0; i < batch_size; ++i) {
 167     while (Skip()) {
 168       Next();
 169     }
 170     for (int j = 0; j < this->layer_param_.top_size(); ++j) {
 171       int data_dim = top[j]->count() / top[j]->shape(0);
 172       caffe_copy(data_dim,
 173           &hdf_blobs_[j]->cpu_data()[data_permutation_[current_row_]
 174             * data_dim], &top[j]->mutable_cpu_data()[i * data_dim]);
 175     }
 176     Next();
 177   }
 178 }
 179
 180 #ifdef CPU_ONLY
 181 STUB_GPU_FORWARD(HDF5DataLayer, Forward);
 182 #endif
 183
 184 INSTANTIATE_CLASS(HDF5DataLayer);
 185 REGISTER_LAYER_CLASS(HDF5Data);
 186
 187 }  // namespace caffe
 188 #endif  // USE_HDF5