data_reader.hpp

Go to the documentation of this file.

// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); you
// may not use this file except in compliance with the License.  You may
// obtain a copy of the License at:
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the license.
//
// lbann_data_reader .hpp - Input data base class for training, testing

#ifndef LBANN_DATA_READER_HPP
#define LBANN_DATA_READER_HPP

#include "lbann/base.hpp"
#include "lbann/data_readers/metadata.hpp"
#include "lbann/data_readers/utils/input_data_type.hpp"
#include "lbann/io/file_io.hpp"
#include "lbann/transforms/transform_pipeline.hpp"
#include "lbann/utils/argument_parser.hpp"
#include "lbann/utils/options.hpp"
#include "lbann/utils/random_number_generators.hpp"

#include <algorithm>
#include <cassert>
#include <string>
#include <unistd.h>
#include <unordered_set>
#include <vector>

// Forward-declare Conduit nodes.
namespace conduit {
class Node;
}

#define NOT_IMPLEMENTED(n)                                                     \
  {                                                                            \
    std::stringstream s;                                                       \
    s << "the method " << n << " has not been implemented";                    \
    throw lbann_exception(s.str());                                            \
  }

namespace lbann {

// Forward declarations
class Layer;
class data_store_conduit;
class thread_pool;
class trainer;
class persist;

class generic_data_reader
{
public:
  using unused_index_map_t = std::map<execution_mode, std::vector<int>>;

  generic_data_reader(bool shuffle = true)
    : m_verbose(global_argument_parser().get<bool>(LBANN_OPTION_VERBOSE)),
      m_data_store(nullptr),
      m_comm(nullptr),
      m_mini_batch_size(0),
      m_current_pos(0),
      m_stride_to_next_mini_batch(0),
      m_base_offset(0),
      m_sample_stride(1),
      m_iteration_stride(1),
      m_last_mini_batch_size(0),
      m_stride_to_last_mini_batch(0),
      m_reset_mini_batch_index(0),
      m_loaded_mini_batch_idx(0),
      m_current_mini_batch_idx(0),
      m_num_iterations_per_epoch(0),
      m_num_parallel_readers(0),
      m_max_files_to_load(0),
      m_file_dir(""),
      m_data_sample_list(""),
      m_data_fn(""),
      m_label_fn(""),
      m_shuffle(shuffle),
      m_absolute_sample_count(0),
      m_use_fraction(1.0),
      m_gan_labelling(false), // default, not GAN
      m_gan_label_value(
        0), // If GAN, default for fake label, discriminator model
      m_io_thread_pool(nullptr),
      m_keep_sample_order(false),
      m_issue_warning(true)
  {
    // By default only support fetching input samples
    m_supported_input_types[INPUT_DATA_TYPE_SAMPLES] = true;
  }
  generic_data_reader(const generic_data_reader&) = default;
  generic_data_reader& operator=(const generic_data_reader&) = default;

  virtual ~generic_data_reader();
  virtual generic_data_reader* copy() const = 0;

  template <class Archive>
  void serialize(Archive& ar);

  void set_comm(lbann_comm* comm) { m_comm = comm; }

  lbann_comm* get_comm() const { return m_comm; }

  virtual bool has_conduit_output() { return false; }

  // These non-virtual methods are used to specify where data is, how much to
  // load, etc.

  void set_file_dir(std::string s);

  void set_local_file_dir(std::string s);

  void set_max_files_to_load(size_t n) { m_max_files_to_load = n; }

  std::string get_file_dir() const;

  std::string get_local_file_dir() const;

  void set_data_sample_list(std::string s);

  std::string get_data_sample_list() const;

  void keep_sample_order(bool same_order = false);

  void set_data_filename(std::string s);

  std::string get_data_filename() const;

  void set_label_filename(std::string s);

  std::string get_label_filename() const;

  void set_shuffle(bool b) { m_shuffle = b; }

  bool is_shuffled() const { return m_shuffle; }

  void set_shuffled_indices(const std::vector<int>& indices)
  {
    m_shuffled_indices = indices;
  }

  const std::vector<int>& get_shuffled_indices() const
  {
    return m_shuffled_indices;
  }

  void set_first_n(int n);

  void set_absolute_sample_count(size_t s);

  void set_use_fraction(double s);

  virtual void set_execution_mode_split_fraction(execution_mode m, double s);

  virtual void set_role(std::string role);

  std::string get_role() const { return m_role; }

  virtual void load() = 0;

  virtual void setup(int num_io_threads,
                     observer_ptr<thread_pool> io_thread_pool);

  virtual std::string get_type() const = 0;

  int fetch(std::map<data_field_type, CPUMat*>& input_buffers,
            El::Matrix<El::Int>& indices_fetched,
            size_t mb_size);

  int fetch(std::vector<conduit::Node>& samples,
            El::Matrix<El::Int>& indices_fetched,
            size_t mb_size);

  virtual bool has_data_field(data_field_type data_field) const
  {
    if (m_supported_input_types.find(data_field) !=
        m_supported_input_types.end()) {
      return m_supported_input_types.at(data_field);
    }
    else {
      return false;
    }
  }

  virtual bool has_labels() const
  {
    return has_data_field(INPUT_DATA_TYPE_LABELS);
  }
  virtual bool has_responses() const
  {
    return has_data_field(INPUT_DATA_TYPE_RESPONSES);
  }

  void set_has_data_field(data_field_type const data_field, const bool b)
  {
    m_supported_input_types[data_field] = b;
  }

  virtual void set_has_labels(const bool b)
  {
    m_supported_input_types[INPUT_DATA_TYPE_LABELS] = b;
  }
  virtual void set_has_responses(const bool b)
  {
    m_supported_input_types[INPUT_DATA_TYPE_RESPONSES] = b;
  }

  void start_data_store_mini_batch_exchange();
  void finish_data_store_mini_batch_exchange();

  virtual bool update(bool is_active_reader);

  virtual int get_num_labels() const { return 0; }
  virtual int get_num_responses() const { return 1; }
  virtual int get_linearized_data_size() const { return 0; }
  virtual int get_linearized_label_size() const { return 0; }
  virtual int get_linearized_response_size() const { return 1; }
  virtual int get_linearized_size(data_field_type const& data_field) const;

  virtual const std::vector<El::Int> get_data_dims() const
  {
    return std::vector<El::Int>(0);
  }

  virtual std::vector<El::Int>
  get_slice_points(const slice_points_mode var_category, bool& is_supported)
  {
    is_supported = false;
    return {};
  }

  virtual bool position_valid() const
  {
    return (m_current_pos < get_num_data());
  }
  virtual bool position_is_overrun() const
  {
    int end_pos = (int)m_shuffled_indices.size();
    return (m_current_pos >= end_pos &&
            (m_current_pos - end_pos) < m_comm->get_procs_per_trainer());
  }
  bool at_new_epoch() const
  {
    return ((m_loaded_mini_batch_idx == m_reset_mini_batch_index) &&
            (m_current_mini_batch_idx == 0));
  }
  void set_mini_batch_size(const int s);
  int get_mini_batch_size() const { return m_mini_batch_size; }
  int get_loaded_mini_batch_size() const;
  int get_current_mini_batch_size() const;
  int get_mini_batch_max() const { return m_mini_batch_size; }
  void set_stride_to_next_mini_batch(const int s)
  {
    m_stride_to_next_mini_batch = s;
  }
  int get_stride_to_next_mini_batch() const
  {
    return m_stride_to_next_mini_batch;
  }
  void set_sample_stride(const int s) { m_sample_stride = s; }
  int get_sample_stride() const { return m_sample_stride; }
  void set_iteration_stride(const int s) { m_iteration_stride = s; }
  int get_iteration_stride() const { return m_iteration_stride; }
  virtual void set_base_offset(const int s) { m_base_offset = s; }
  int get_base_offset() const { return m_base_offset; }
  void set_last_mini_batch_size(const int s) { m_last_mini_batch_size = s; }
  int get_last_mini_batch_size() const { return m_last_mini_batch_size; }
  void set_stride_to_last_mini_batch(const int s)
  {
    m_stride_to_last_mini_batch = s;
  }
  int get_stride_to_last_mini_batch() const
  {
    return m_stride_to_last_mini_batch;
  }
  void set_num_parallel_readers(const int s) { m_num_parallel_readers = s; }
  int get_num_parallel_readers() const { return m_num_parallel_readers; }
  virtual void set_reset_mini_batch_index(const int s)
  {
    m_reset_mini_batch_index = s;
  }
  int get_reset_mini_batch_index() const { return m_reset_mini_batch_index; }
  int get_loaded_mini_batch_index() const { return m_loaded_mini_batch_idx; }
  int get_current_mini_batch_index() const { return m_current_mini_batch_idx; }
  void set_initial_position()
  {
    m_current_pos = m_base_offset;
    m_loaded_mini_batch_idx = m_reset_mini_batch_index;
    m_current_mini_batch_idx = 0;
  }
  int get_position() const { return m_current_pos; }
  int get_next_position() const;
  int* get_indices() { return &m_shuffled_indices[0]; }
  virtual int get_num_data() const { return (int)m_shuffled_indices.size(); }
  int get_num_unused_data(execution_mode m) const;

  int* get_unused_data(execution_mode m);

  const std::vector<int>& get_unused_indices(execution_mode m);

  void set_num_iterations_per_epoch(int num_iterations_per_epoch)
  {
    m_num_iterations_per_epoch =
      num_iterations_per_epoch; 
  }
  int get_num_iterations_per_epoch() const
  {
    return m_num_iterations_per_epoch; 
  }

  int get_current_step_in_epoch() const { return m_current_mini_batch_idx; }

  void resize_shuffled_indices();

  void select_subset_of_data();

  virtual void use_unused_index_set(execution_mode m);

  virtual bool has_list_per_model() const { return false; }
  virtual bool has_list_per_trainer() const { return false; }

  bool save_to_checkpoint_shared(persist& p, execution_mode mode);

  bool load_from_checkpoint_shared(persist& p, execution_mode mode);

  bool save_to_checkpoint_distributed(persist& p, execution_mode mode);

  bool load_from_checkpoint_distributed(persist& p, execution_mode mode);

  const data_store_conduit& get_data_store() const;

  data_store_conduit& get_data_store();

  data_store_conduit* get_data_store_ptr() const { return m_data_store; }

  void setup_data_store(int mini_batch_size);

  void instantiate_data_store();

  virtual void preload_data_store();

  void set_gan_labelling(bool has_gan_labelling)
  {
    m_gan_labelling = has_gan_labelling;
  }
  void set_gan_label_value(int gan_label_value)
  {
    m_gan_label_value = gan_label_value;
  }

  void set_data_store(data_store_conduit* g);

  virtual bool data_store_active() const;

  virtual bool priming_data_store() const;

  virtual void post_update() {}

  void set_transform_pipeline(transform::transform_pipeline&& tp)
  {
    m_transform_pipeline = std::move(tp);
  }

#ifdef LBANN_HAS_DISTCONV

  virtual bool is_tensor_shuffle_required() const { return true; }
#endif // LBANN_HAS_DISTCONV

protected:
  bool m_verbose = false;

  // For use with conduit when samples are corrupt.
  mutable std::unordered_set<int> m_using_random_node;

  size_t get_absolute_sample_count() const;

  double get_use_fraction() const;

  double get_execution_mode_split_fraction(execution_mode m) const;

  data_store_conduit* m_data_store;

  lbann_comm* m_comm;

  virtual bool
  fetch_data_block(std::map<data_field_type, CPUMat*>& input_buffers,
                   El::Int block_offset,
                   El::Int block_stride,
                   El::Int mb_size,
                   El::Matrix<El::Int>& indices_fetched);

  bool fetch_data_block_conduit(std::vector<conduit::Node>& samples,
                                El::Int block_offset,
                                El::Int block_stride,
                                El::Int mb_size,
                                El::Matrix<El::Int>& indices_fetched);

  virtual bool fetch_data_field(data_field_type data_field,
                                CPUMat& Y,
                                int data_id,
                                int mb_idx)
  {
    NOT_IMPLEMENTED("fetch_data_field");
    return false;
  }

  virtual bool fetch_conduit_node(conduit::Node& sample, int data_id)
  {
    NOT_IMPLEMENTED("fetch_conduit_node");
    return false;
  }

  virtual bool fetch_datum(CPUMat& X, int data_id, int mb_idx)
  {
    NOT_IMPLEMENTED("fetch_dataum");
    return false;
  }

  virtual bool fetch_label(CPUMat& Y, int data_id, int mb_idx)
  {
    NOT_IMPLEMENTED("fetch_label");
    return false;
  }

  virtual bool fetch_response(CPUMat& Y, int data_id, int mb_idx)
  {
    NOT_IMPLEMENTED("fetch_response");
    return false;
  }

  inline CPUMat create_datum_view(CPUMat& X, const int mb_idx)
  {
    return El::View(X, El::IR(0, X.Height()), El::IR(mb_idx, mb_idx + 1));
  }

  virtual void preprocess_data_source(int tid){};
  virtual void postprocess_data_source(int tid){};

  virtual void shuffle_indices();
  virtual void shuffle_indices(rng_gen& gen);

public:
  int m_mini_batch_size;
  int m_current_pos;
  int m_stride_to_next_mini_batch;
  int m_base_offset;
  int m_sample_stride;
  int m_iteration_stride;

  std::vector<int> m_shuffled_indices;
  unused_index_map_t m_unused_indices;

  int m_last_mini_batch_size;
  int m_stride_to_last_mini_batch;
  int m_reset_mini_batch_index;
  int m_loaded_mini_batch_idx;
  int m_current_mini_batch_idx;
  int
    m_num_iterations_per_epoch; 

  int m_num_parallel_readers; 

  size_t m_max_files_to_load;
  std::string m_file_dir;
  std::string m_local_file_dir;
  std::string m_data_sample_list;
  std::string m_data_fn;
  std::string m_label_fn;
  bool m_shuffle;
  size_t m_absolute_sample_count;
  std::map<execution_mode, double> m_execution_mode_split_fraction;
  double m_use_fraction;
  int m_first_n;
  std::string m_role;

  void print_get_methods(const std::string filename);

  size_t get_num_indices_to_use() const;

  friend class data_reader_merge_features;
  friend class data_reader_merge_samples;

  void set_use_data_store(bool s) { m_use_data_store = s; }

private:
  virtual void do_preload_data_store(); // Throws: "Not implemented."

protected:
  bool m_use_data_store = false;

  std::map<data_field_type, bool> m_supported_input_types;

  // var to support GAN
  bool m_gan_labelling; // boolean flag of whether its GAN binary label, default
                        // is false
  int m_gan_label_value; // zero(0) or 1 label value for discriminator, default
                         // is 0

  observer_ptr<thread_pool> m_io_thread_pool;

  bool m_keep_sample_order;

  transform::transform_pipeline m_transform_pipeline;

  bool m_issue_warning;

  void error_check_counts() const;
};

template <typename T>
inline void set_minibatch_item(Mat& M,
                               const int mb_idx,
                               const T* const ptr,
                               const size_t count)
{
  if ((count > 0u) && (ptr == nullptr)) {
    throw lbann_exception(std::string{} + __FILE__ + " " +
                          std::to_string(__LINE__) +
                          " :: attempt to dereference a nullptr ");
  }
  for (size_t i = 0u; i < count; ++i) {
    M.Set(static_cast<El::Int>(i),
          static_cast<El::Int>(mb_idx),
          static_cast<DataType>(ptr[i]));
  }
}

} // namespace lbann

#endif // LBANN_DATA_READER_HPP