data_reader_HDF5.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 LBANN_DATA_READER_HDF5_REVISED_HPP
#define LBANN_DATA_READER_HDF5_REVISED_HPP

#include "lbann/data_readers/data_reader.hpp"
#include "lbann/data_readers/data_reader_sample_list.hpp"
#include "lbann/data_readers/sample_list_hdf5.hpp"
#include "lbann/data_store/data_store_conduit.hpp"

#include <set>

// Forward declaration
class DataReaderHDF5WhiteboxTester;

#define HDF5_METADATA_KEY_DIMS "dims"
#define HDF5_METADATA_KEY_CHANNELS "channels"
#define HDF5_METADATA_KEY_ORDERING "ordering"
#define HDF5_METADATA_KEY_SCALE "scale"
#define HDF5_METADATA_KEY_BIAS "bias"
#define HDF5_METADATA_KEY_LAYOUT "layout"
#define HDF5_METADATA_KEY_TRANSPOSE "transpose"
#define HDF5_METADATA_KEY_COERCE "coerce"
#define HDF5_METADATA_KEY_PACK "pack"

#define HDF5_METADATA_VALUE_COERCE_FLOAT "float"
#define HDF5_METADATA_VALUE_COERCE_DOUBLE "double"
#define HDF5_METADATA_VALUE_COERCE_FLOAT64 "float64"
#define HDF5_METADATA_VALUE_COERCE_FLOAT16 "float16"
#define HDF5_METADATA_VALUE_LAYOUT_CHW "chw"
#define HDF5_METADATA_VALUE_LAYOUT_HWC "hwc"
#define HDF5_METADATA_VALUE_LAYOUT_CDHW "cdhw"
#define HDF5_METADATA_VALUE_LAYOUT_DHWC "dhwc"

namespace lbann {

bool is_hdf5_metadata_key_valid(std::string const& key);
bool is_hdf5_field_channels_last(conduit::Node const& field);
bool does_hdf5_field_require_repack_to_channels_first(
  conduit::Node const& metadata);
std::string conduit_to_string(conduit::Node const& field);

static std::set<std::string> const hdf5_metadata_valid_keys = {
  HDF5_METADATA_KEY_DIMS,
  HDF5_METADATA_KEY_CHANNELS,
  HDF5_METADATA_KEY_ORDERING,
  HDF5_METADATA_KEY_SCALE,
  HDF5_METADATA_KEY_BIAS,
  HDF5_METADATA_KEY_LAYOUT,
  HDF5_METADATA_KEY_TRANSPOSE,
  HDF5_METADATA_KEY_COERCE,
  HDF5_METADATA_KEY_PACK,
};

class hdf5_data_reader
  : public data_reader_sample_list<sample_list_hdf5<std::string>>
{
public:
  hdf5_data_reader(bool shuffle = true);
  hdf5_data_reader(const hdf5_data_reader&);
  hdf5_data_reader& operator=(const hdf5_data_reader&);
  hdf5_data_reader* copy() const override
  {
    return new hdf5_data_reader(*this);
  }
  void copy_members(const hdf5_data_reader& rhs);
  ~hdf5_data_reader() override;

  bool has_conduit_output() override { return true; }

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

  void print_metadata(std::ostream& os = std::cout);

  void load() override;

  bool fetch_conduit_node(conduit::Node& sample, int data_id) override;

  void set_experiment_schema_filename(std::string fn)
  {
    m_experiment_schema_filename = fn;
  }

  const std::string& get_experiment_schema_filename()
  {
    return m_experiment_schema_filename;
  }

  void set_data_schema_filename(std::string fn) { m_data_schema_filename = fn; }

  const std::string& get_data_schema_filename()
  {
    return m_data_schema_filename;
  }

  const std::vector<El::Int> get_data_dims() const override
  {
    return get_data_dims(INPUT_DATA_TYPE_SAMPLES);
  }

  int get_linearized_data_size() const override
  {
    return get_linearized_size(INPUT_DATA_TYPE_SAMPLES);
  }

  int get_linearized_response_size() const override
  {
    return get_linearized_size(INPUT_DATA_TYPE_RESPONSES);
  }

  int get_linearized_label_size() const override
  {
    return get_linearized_size(INPUT_DATA_TYPE_LABELS);
  }

  int get_num_labels() const override { return get_linearized_label_size(); }

  int get_num_responses() const override
  {
    return get_linearized_response_size();
  }

  conduit::Node get_experiment_schema() const { return m_experiment_schema; }
  conduit::Node get_data_schema() const { return m_data_schema; }
  void set_experiment_schema(const conduit::Node& s);
  void set_data_schema(const conduit::Node& s);
  std::unordered_map<std::string, conduit::Node> get_node_map() const
  {
    return m_useme_node_map;
  }

  void adjust_metadata(conduit::Node* root);

private:
  std::unordered_map<std::string, std::vector<El::Int>> m_data_dims_lookup_table;

  std::unordered_map<std::string, int> m_linearized_size_lookup_table;

  std::string m_experiment_schema_filename;

  std::string m_data_schema_filename;

  // to set to false, use the cmd line flag: --keep_packed_fields
  // I don't know a use case for keeping an original (possibly coerced)
  // field if it's been packed, but someone else might ...
  // note that setting to 'false' invokes both a memory and communication
  // penalty
  bool m_delete_packed_fields = true;

  struct PackingGroup
  {
    std::string group_name;
    conduit::index_t data_type;
    size_t n_elts;
    std::vector<std::string> names;
    std::vector<size_t> sizes;
    std::vector<conduit::index_t> data_types;
  };

  std::unordered_map<std::string, PackingGroup> m_packing_groups;

  const std::string s_metadata_node_name = "metadata";

  const std::string s_composite_node = "composite_node";

  std::unordered_map<std::string, conduit::Node*> m_useme_node_map_ptrs;
  std::unordered_map<std::string, conduit::Node> m_useme_node_map;

  conduit::Node m_experiment_schema;

  conduit::Node m_data_schema;

  std::unordered_map<std::string, conduit::Node*> m_data_map;

  std::unordered_set<std::string> m_add_to_map;

  //=========================================================================
  // methods follow
  //=========================================================================

  const std::vector<El::Int> get_data_dims(std::string name = "") const;

  int get_linearized_size(data_field_type const& data_field) const override;

  void load_sample_schema(conduit::Schema& s);

  void parse_schemas();

  void get_schema_ptrs(
    conduit::Node* starting_node,
    std::unordered_map<std::string, conduit::Node*>& schema_name_map);

  void get_leaves(conduit::Node* node_in,
                  std::unordered_map<std::string, conduit::Node*>& leaves_out);

  void
  get_leaves_multi(conduit::Node* node_in,
                   std::unordered_map<std::string, conduit::Node*>& leaves_out);

  void do_preload_data_store() override;

  void load_sample(conduit::Node& node,
                   hid_t file_handle,
                   const std::string& sample_name,
                   bool ignore_failure = false);

  void load_sample_from_sample_list(conduit::Node& node,
                                    size_t index,
                                    bool ignore_failure = false);

  void pack_data(conduit::Node& node_in_out);

  void load_schema(std::string filename, conduit::Node& schema);

  void pack(conduit::Node& node, size_t index);

  conduit::Node merge_metadata_nodes(const conduit::Node* node_A,
                                     const conduit::Node* node_B);

  void build_packing_map(conduit::Node& node);

  void repack_image(conduit::Node& node,
                    const std::string& path,
                    const conduit::Node& metadata);

  void coerce(const conduit::Node& metadata,
              hid_t file_handle,
              const std::string& original_path,
              const std::string& new_pathname,
              conduit::Node& node);

  void normalize(conduit::Node& node,
                 const std::string& path,
                 const conduit::Node& metadata);

  void construct_linearized_size_lookup_tables();
  void construct_linearized_size_lookup_tables(conduit::Node& node);

  void test_that_all_nodes_contain_metadata(conduit::Node& node);

  bool get_delete_packed_fields() { return m_delete_packed_fields; }
  void set_delete_packed_fields(bool flag) { m_delete_packed_fields = flag; }

  //=========================================================================
  // template declarations follow
  //=========================================================================

  template <typename T>
  void pack(std::string const& group_name, conduit::Node& node, size_t index);

  bool is_composite_node(const conduit::Node& node) const;

  // Designate a whitebox testing friend
  friend class ::DataReaderHDF5WhiteboxTester;

}; // END: class hdf5_data_reader

} // namespace lbann

#endif // LBANN_DATA_READER_HDF5_REVISED_HPP