src/caffe/layers/hdf5_data_layer.cu

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