data_reader_jag_conduit.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.

#ifndef _DATA_READER_JAG_CONDUIT_HPP_
#define _DATA_READER_JAG_CONDUIT_HPP_

#include "lbann_config.hpp"

#include "conduit/conduit.hpp"
#include "hdf5.h"
#include "lbann/data_readers/data_reader.hpp"
#include <map>
#include <memory>
#include <set>
#include <string>
#include <unordered_map>

// #define _USE_IO_HANDLE_
#ifdef _USE_IO_HANDLE_
#include "lbann/data_readers/sample_list_conduit_io_handle.hpp"
#else
#include "lbann/data_readers/sample_list_hdf5.hpp"
#endif

namespace lbann {

class data_store_conduit;

class data_reader_jag_conduit : public generic_data_reader
{
public:
  using ch_t = float; 
  using conduit_ch_t =
    conduit::float32_array; 
  using scalar_t = double;  
  using input_t = double;   
  using sample_locator_t = std::pair<std::string, hid_t>;
  using sample_map_t = std::vector<sample_locator_t>; 
  using sample_name_t = std::string;
#ifdef _USE_IO_HANDLE_
  using sample_list_t = sample_list_conduit_io_handle<sample_name_t>;
#else
  using sample_list_t = sample_list_hdf5<sample_name_t>;
#endif
  using file_handle_t = sample_list_t::file_handle_t;
  using sample_file_id_t = sample_list_t::sample_file_id_t;
  using sample_t = std::pair<sample_file_id_t, sample_name_t>;
  // using sample_t = sample_list_t::sample_t;
  using linear_transform_t = std::pair<double, double>;

  enum variable_t
  {
    Undefined = 0,
    JAG_Image,
    JAG_Scalar,
    JAG_Input
  };
  using TypeID = conduit::DataType::TypeID;

  using prefix_t = std::pair<std::string, size_t>;

  data_reader_jag_conduit(bool shuffle = true);
  data_reader_jag_conduit(const data_reader_jag_conduit&);
  data_reader_jag_conduit& operator=(const data_reader_jag_conduit&);
  ~data_reader_jag_conduit() override;
  data_reader_jag_conduit* copy() const override
  {
    return new data_reader_jag_conduit(*this);
  }

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

  std::string get_type() const override { return "data_reader_jag_conduit"; }

  void set_independent_variable_type(
    const std::vector<std::vector<variable_t>>& independent);
  void set_dependent_variable_type(
    const std::vector<std::vector<variable_t>>& dependent);

  std::vector<variable_t> get_independent_variable_type() const;
  std::vector<variable_t> get_dependent_variable_type() const;

  void set_output_scalar_prefix(const std::string& prefix)
  {
    m_output_scalar_prefix = prefix;
  }
  void set_output_image_prefix(const std::string& prefix)
  {
    m_output_image_prefix = prefix;
  }
  void set_input_prefix(const std::string& prefix) { m_input_prefix = prefix; }

  void set_image_dims(const int width, const int height, const int ch = 1);
  void set_image_choices(const std::vector<std::string> image_keys);
  const std::vector<std::string>& get_image_choices() const;

  void add_scalar_filter(const std::string& key);
  void add_scalar_prefix_filter(const prefix_t& p);
  void add_input_filter(const std::string& key);
  void add_input_prefix_filter(const prefix_t& p);

  void set_scalar_choices(const std::vector<std::string>& keys);
  void set_all_scalar_choices();
  const std::vector<std::string>& get_scalar_choices() const;

  void set_input_choices(const std::vector<std::string>& keys);
  void set_all_input_choices();
  const std::vector<std::string>& get_input_choices() const;

#ifndef _JAG_OFFLINE_TOOL_MODE_
  void load() override;

  void set_local_id(const std::string role);
  int get_local_id(const std::string role) const;
  // void set_open_hdf5_files(std::shared_ptr<hdf5_file_handles>& f);
  // /// Get the set of open hdf5 data files
  // std::shared_ptr<hdf5_file_handles>& get_open_hdf5_files();
#else
  void check_image_data();
#endif // _JAG_OFFLINE_TOOL_MODE_

  void set_list_per_trainer(bool flag) { m_list_per_trainer = flag; };
  void set_list_per_model(bool flag) { m_list_per_model = flag; };

  bool has_list_per_model() const override { return m_list_per_model; }
  bool has_list_per_trainer() const override { return m_list_per_trainer; }

  unsigned int get_num_img_srcs() const;
  size_t get_linearized_image_size() const;
  size_t get_linearized_1ch_image_size() const;
  size_t get_linearized_scalar_size() const;
  size_t get_linearized_input_size() const;

  int get_linearized_data_size() const override;
  int get_linearized_response_size() const override;
  std::vector<size_t> get_linearized_data_sizes() const;
  std::vector<size_t> get_linearized_response_sizes() const;

  const std::vector<El::Int> get_data_dims() const override;

  int get_num_data() const override;
  int get_num_labels() const override;
  int get_linearized_label_size() const override;

  std::vector<El::Int> get_slice_points(const slice_points_mode var_category,
                                        bool& is_supported) override;

  void set_split_image_channels();
  void unset_split_image_channels();
  bool check_split_image_channels() const;

  std::string get_description() const;

  std::vector<scalar_t> get_scalars(const size_t i,
                                    conduit::Node& sample) const;

  std::vector<input_t> get_inputs(const size_t i, conduit::Node& sample) const;

  template <typename S>
  static size_t
  add_val(const std::string key, const conduit::Node& n, std::vector<S>& vals);

  void setup_data_store(int mini_batch_size);

  static std::string to_string(const variable_t t);

  void print_schema(const size_t i) const;

  void clear_image_normalization_params();
  void clear_scalar_normalization_params();
  void clear_input_normalization_params();
  void add_image_normalization_param(const linear_transform_t& t);
  void add_scalar_normalization_param(const linear_transform_t& t);
  void add_input_normalization_param(const linear_transform_t& t);

protected:
  void do_preload_data_store() override;

  virtual void set_defaults();
  virtual void copy_members(const data_reader_jag_conduit& rhs);

  void add_independent_variable_type(const variable_t independent);
  void add_dependent_variable_type(const variable_t dependent);

  bool filter(const std::set<std::string>& key_filter,
              const std::vector<prefix_t>& prefix_filter,
              const std::string& name) const;

  using generic_data_reader::get_linearized_size;
  size_t get_linearized_size(const variable_t t) const;
  const std::vector<int> get_dims(const variable_t t) const;
  std::vector<El::Int> get_slice_points_impl(
    const std::vector<std::vector<data_reader_jag_conduit::variable_t>>& var)
    const;
  std::vector<El::Int> get_slice_points_independent() const;
  std::vector<El::Int> get_slice_points_dependent() const;

  static std::string to_string(const std::vector<variable_t>& vec);
  static std::string to_string(const std::vector<std::vector<variable_t>>& vec);

  virtual std::vector<CPUMat>
  create_datum_views(CPUMat& X,
                     const std::vector<size_t>& sizes,
                     const int mb_idx) const;


  bool fetch(CPUMat& X,
             int data_id,
             conduit::Node& sample,
             int mb_idx,
             int tid,
             const variable_t vt,
             const std::string tag);
  bool fetch_datum(CPUMat& X, int data_id, int mb_idx) override;
  bool fetch_response(CPUMat& Y, int data_id, int mb_idx) override;
  bool fetch_label(CPUMat& X, int data_id, int mb_idx) override;

#ifndef _JAG_OFFLINE_TOOL_MODE_
  using generic_data_reader::shuffle_indices;

  void shuffle_indices(rng_gen& gen) override;

  int compute_max_num_parallel_readers();
  bool check_num_parallel_readers(long data_set_size);
  void sample_schema_check(const bool check_data);
  void load_list_of_samples(const std::string filename);
  void
  load_list_of_samples_from_archive(const std::string& sample_list_archive);

  void check_image_data();
#endif // _JAG_OFFLINE_TOOL_MODE_

#if 0
  hid_t open_conduit_file(const std::string& conduit_file_path);

  void open_all_conduit_files();
#endif

  void set_linearized_image_size();
  void check_scalar_keys();
  void check_input_keys();

  static bool check_non_numeric(const std::string key);

  bool has_path(const file_handle_t& h, const std::string& path) const;
  void read_node(const file_handle_t& h,
                 const std::string& path,
                 conduit::Node& n) const;

  static const conduit::Node& get_conduit_node(const conduit::Node& n_base,
                                               const std::string key);
  bool load_conduit_node(const size_t i,
                         const std::string& key,
                         conduit::Node& node) const;
  bool has_conduit_path(const size_t i, const std::string& key) const;

  std::vector<std::vector<DataType>>
  get_image_data(const size_t i, conduit::Node& sample) const;

  bool data_store_active() const override
  {
    bool flag = generic_data_reader::data_store_active();
    return (m_data_store != nullptr && flag);
  }

  bool priming_data_store() const override
  {
    bool flag = generic_data_reader::priming_data_store();
    return (m_data_store != nullptr && flag);
  }

protected:
  std::vector<variable_t> m_independent;
  std::vector<std::vector<variable_t>> m_independent_groups;
  std::vector<variable_t> m_dependent;
  std::vector<std::vector<variable_t>> m_dependent_groups;

  int m_image_width;                  
  int m_image_height;                 
  int m_image_num_channels;           
  size_t m_image_linearized_size;     
  size_t m_1ch_image_linearized_size; 
  unsigned int m_num_img_srcs;        
  bool m_split_channels; 

  bool m_is_data_loaded;

  int m_num_labels; 

  std::string m_output_scalar_prefix;
  std::string m_output_image_prefix;
  std::string m_input_prefix;

  std::vector<std::string> m_emi_image_keys;
  std::vector<std::string> m_scalar_keys;
  std::vector<std::string> m_input_keys;

  static const std::set<std::string> non_numeric_vars;

  bool m_uniform_input_type;

  std::set<std::string> m_scalar_filter;
  std::vector<prefix_t> m_scalar_prefix_filter;
  std::set<std::string> m_input_filter;
  std::vector<prefix_t> m_input_prefix_filter;

  static std::unordered_map<std::string, int> m_num_local_readers;
  int m_local_reader_id;

  CPUMat m_data_cache;
  CPUMat m_response_cache;
  CPUMat m_label_cache;
  int m_cached_data_mb_size;
  int m_cached_response_mb_size;
  int m_cached_label_mb_size;

  std::vector<linear_transform_t> m_image_normalization_params;
  std::vector<linear_transform_t> m_scalar_normalization_params;
  std::vector<linear_transform_t> m_input_normalization_params;

  typedef std::pair<std::string, std::string> conduit_sample;
  sample_list_t m_sample_list;
  bool m_list_per_trainer;
  bool m_list_per_model;

  void preload_helper(const hid_t& h,
                      const std::string& sample_name,
                      const std::string& field_name,
                      int data_id,
                      conduit::Node& node);
};

template <typename TN, conduit::DataType::TypeID TC>
struct is_same : std::false_type
{
};

#define _LBANN_CONDUIT_DTYPE_INSTANTIATION_(TN, TC)                            \
  template <>                                                                  \
  struct is_same<TN, TC> : std::true_type                                      \
  {                                                                            \
  }

_LBANN_CONDUIT_DTYPE_INSTANTIATION_(int8_t, conduit::DataType::INT8_ID);
_LBANN_CONDUIT_DTYPE_INSTANTIATION_(int16_t, conduit::DataType::INT16_ID);
_LBANN_CONDUIT_DTYPE_INSTANTIATION_(int32_t, conduit::DataType::INT32_ID);
_LBANN_CONDUIT_DTYPE_INSTANTIATION_(int64_t, conduit::DataType::INT64_ID);
_LBANN_CONDUIT_DTYPE_INSTANTIATION_(uint8_t, conduit::DataType::UINT8_ID);
_LBANN_CONDUIT_DTYPE_INSTANTIATION_(uint16_t, conduit::DataType::UINT16_ID);
_LBANN_CONDUIT_DTYPE_INSTANTIATION_(uint32_t, conduit::DataType::UINT32_ID);
_LBANN_CONDUIT_DTYPE_INSTANTIATION_(uint64_t, conduit::DataType::UINT64_ID);
_LBANN_CONDUIT_DTYPE_INSTANTIATION_(float, conduit::DataType::FLOAT32_ID);
_LBANN_CONDUIT_DTYPE_INSTANTIATION_(double, conduit::DataType::FLOAT64_ID);
_LBANN_CONDUIT_DTYPE_INSTANTIATION_(char*, conduit::DataType::CHAR8_STR_ID);

#undef _LBANN_CONDUIT_DTYPE_INSTANTIATION_

template <typename TN>
inline bool is_same_type(const conduit::DataType::TypeID dt)
{
  switch (dt) {
  case conduit::DataType::INT8_ID:
    return is_same<TN, conduit::DataType::INT8_ID>::value;
  case conduit::DataType::INT16_ID:
    return is_same<TN, conduit::DataType::INT16_ID>::value;
  case conduit::DataType::INT32_ID:
    return is_same<TN, conduit::DataType::INT32_ID>::value;
  case conduit::DataType::INT64_ID:
    return is_same<TN, conduit::DataType::INT64_ID>::value;
  case conduit::DataType::UINT8_ID:
    return is_same<TN, conduit::DataType::UINT8_ID>::value;
  case conduit::DataType::UINT16_ID:
    return is_same<TN, conduit::DataType::UINT16_ID>::value;
  case conduit::DataType::UINT32_ID:
    return is_same<TN, conduit::DataType::UINT32_ID>::value;
  case conduit::DataType::UINT64_ID:
    return is_same<TN, conduit::DataType::UINT64_ID>::value;
  case conduit::DataType::FLOAT32_ID:
    return is_same<TN, conduit::DataType::FLOAT32_ID>::value;
  case conduit::DataType::FLOAT64_ID:
    return is_same<TN, conduit::DataType::FLOAT64_ID>::value;
  case conduit::DataType::CHAR8_STR_ID:
    return is_same<TN, conduit::DataType::CHAR8_STR_ID>::value;
  default:
    return false;
  }
  return false;
}

template <typename S>
inline size_t data_reader_jag_conduit::add_val(const std::string key,
                                               const conduit::Node& n,
                                               std::vector<S>& vals)
{
  size_t cnt = 0u;

  switch (n.dtype().id()) {
  case TypeID::OBJECT_ID: {
    // std::cout << "O " << n.path() << std::endl;
    if (check_non_numeric(key)) {
      return 0u;
    }
    conduit::NodeConstIterator itr = n.children();
    while (itr.has_next()) {
      const conduit::Node& n_child = itr.next();
      cnt += add_val(itr.name(), n_child, vals);
    }
  } break;
  case TypeID::LIST_ID: {
    // std::cout << "L " << n.path() << std::endl;
    if (check_non_numeric(key)) {
      return 0u;
    }
    conduit::NodeConstIterator itr = n.children();
    while (itr.has_next()) {
      const conduit::Node& n_child = itr.next();
      cnt += add_val(itr.name(), n_child, vals);
    }
  } break;
  case TypeID::INT8_ID:
  case TypeID::INT16_ID:
  case TypeID::INT32_ID:
  case TypeID::INT64_ID:
  case TypeID::UINT8_ID:
  case TypeID::UINT16_ID:
  case TypeID::UINT32_ID:
  case TypeID::UINT64_ID:
  case TypeID::FLOAT32_ID:
  case TypeID::FLOAT64_ID:
    cnt = 1u;
    // std::cout << "N " << n.path() << ": " << static_cast<S>(n.to_value()) <<
    // std::endl;
    vals.push_back(static_cast<S>(n.to_value()));
    break;
  case TypeID::CHAR8_STR_ID: {
    // In case of a charater string, the method to convert it to a float number
    // is specific to each key
    if (check_non_numeric(key)) {
      return 0u;
      //} else if (key ==
      //"some_key_with_non_numeric_values_that_can_be_converted_to_numerics_in_a_specific_way")
      //{
    }
    else {
      const char* c_str = n.as_char8_str();
      // make sure that the std::string does not contain null character
      const std::string str =
        ((c_str == nullptr)
           ? std::string()
           : std::string(c_str, n.dtype().number_of_elements()))
          .c_str();

      cnt = 1u;
      const S v = static_cast<S>(atof(str.c_str()));
      vals.push_back(v);
      // std::cout << "S " << n.path() << ": " << str << " => " << vals.back()
      // << std::endl;
    }
  } break;
  case TypeID::EMPTY_ID:
  default:
    std::string err =
      std::string("data_reader_jag_conduit::add_val() : invalid dtype (") +
      n.dtype().name() + ") for " + n.path() + '.';
#if 1
    std::cerr << err << " Skipping for now." << std::endl;
#else
    throw lbann_exception(err);
#endif
  }
  return cnt;
}

} // end of namespace lbann
#endif // _DATA_READER_JAG_CONDUIT_HPP_