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