src/caffe/layers/hdf5_data_layer.cu

   1 /*
   2 TODO:
   3 - only load parts of the file, in accordance with a prototxt param "max_mem"
   4 */
   5
   6 #include <stdint.h>
   7 #include <string>
   8 #include <vector>
   9
  10 #include "hdf5.h"
  11 #include "hdf5_hl.h"
  12
  13 #include "caffe/layer.hpp"
  14 #include "caffe/util/io.hpp"
  15 #include "caffe/vision_layers.hpp"
  16
  17 namespace caffe {
  18
  19 template <typename Dtype>
  20 Dtype HDF5DataLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype>*>& bottom,
  21       vector<Blob<Dtype>*>* top) {
  22   const int batch_size = this->layer_param_.hdf5_data_param().batch_size();
  23   const int data_count = (*top)[0]->count() / (*top)[0]->num();
  24   const int label_data_count = (*top)[1]->count() / (*top)[1]->num();
  25
  26   for (int i = 0; i < batch_size; ++i, ++current_row_) {
  27     if (current_row_ == data_blob_.num()) {
  28       if (num_files_ > 1) {
  29         current_file_ += 1;
  30
  31         if (current_file_ == num_files_) {
  32           current_file_ = 0;
  33           LOG(INFO) << "looping around to first file";
  34         }
  35
  36         LoadHDF5FileData(hdf_filenames_[current_file_].c_str());
  37       }
  38       current_row_ = 0;
  39     }
  40     caffe_copy(data_count,
  41         &data_blob_.cpu_data()[current_row_ * data_count],
  42         &(*top)[0]->mutable_gpu_data()[i * data_count]);
  43     caffe_copy(label_data_count,
  44         &label_blob_.cpu_data()[current_row_ * label_data_count],
  45         &(*top)[1]->mutable_gpu_data()[i * label_data_count]);
  46   }
  47   return Dtype(0.);
  48 }
  49
  50 INSTANTIATE_CLASS(HDF5DataLayer);
  51
  52 }  // namespace caffe