sample_list.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_READERS_SAMPLE_LIST_HPP
#define LBANN_DATA_READERS_SAMPLE_LIST_HPP

#include <functional>
#include <iostream>
#include <string>
#include <vector>

#include "lbann/utils/file_utils.hpp"

namespace lbann {

// Forward Declarations
class lbann_comm;

static const std::string multi_sample_exclusion = "MULTI-SAMPLE_EXCLUSION";
static const std::string multi_sample_inclusion = "MULTI-SAMPLE_INCLUSION";
static const std::string single_sample = "SINGLE-SAMPLE";
static const std::string multi_sample_inclusion_v2 =
  "MULTI-SAMPLE_INCLUSION_V2";
static const std::string conduit_hdf5_exclusion = "CONDUIT_HDF5_EXCLUSION";
static const std::string conduit_hdf5_inclusion = "CONDUIT_HDF5_INCLUSION";

struct sample_list_header
{
  bool m_is_multi_sample;
  bool m_is_exclusive;
  bool m_no_label_header;
  bool m_has_unused_sample_fields;
  size_t m_included_sample_count;
  size_t m_excluded_sample_count;
  size_t m_num_files;
  std::string m_file_dir;
  std::string m_sample_list_name;
  std::string m_label_filename;

  sample_list_header();

  void set_sample_list_type(const std::string& line1);
  void set_sample_count(const std::string& line2);
  void set_data_file_dir(const std::string& line3);
  void set_label_filename(const std::string& line4);

  bool is_multi_sample() const;
  bool is_exclusive() const;
  bool use_label_header() const;
  bool has_unused_sample_fields() const;
  size_t get_sample_count() const;
  size_t get_num_files() const;
  const std::string& get_file_dir() const;
  const std::string& get_sample_list_name() const;
  void set_sample_list_name(const std::string& n);
  const std::string& get_label_filename() const;
  template <class Archive>
  void serialize(Archive& ar);
};

template <typename sample_name_t>
class sample_list
{
public:
  using name_t = sample_name_t;
  using sample_file_id_t = std::size_t;
  using sample_t = std::template pair<sample_file_id_t, sample_name_t>;
  using samples_t = std::template vector<sample_t>;
  using sample_idx_t = typename samples_t::size_type;
  using sample_map_t = std::unordered_map<sample_name_t, sample_idx_t>;
  using file_id_stats_v_t = std::vector<std::string>;

  sample_list();
  virtual ~sample_list();
  sample_list(const sample_list& rhs);
  sample_list& operator=(const sample_list& rhs);
  sample_list& copy(const sample_list& rhs);

  void copy_members(const sample_list& rhs);

  void load(std::istream& istrm, size_t stride = 1, size_t offset = 0);

  void load(const std::string& samplelist_file,
            const lbann_comm& comm,
            bool interleave);
  void load(std::istream& istrm, const lbann_comm& comm, bool interleave);
  void load(const sample_list_header& header,
            std::istream& istrm,
            const lbann_comm& comm,
            bool interleave);

  void load_from_string(const std::string& samplelist,
                        const lbann_comm& comm,
                        bool interleave);

  virtual size_t size() const;

  virtual size_t get_num_files() const;

  bool empty() const;

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

  virtual bool to_string(std::string& sstr) const;

  void write(const std::string filename) const;

  const samples_t& get_list() const;

  const sample_list_header& get_header() const;

  const sample_t& operator[](size_t idx) const;

  virtual const std::string& get_samples_filename(sample_file_id_t id) const;

  const std::string& get_samples_dirname() const;
  const std::string& get_label_filename() const;

  void all_gather_archive(const std::string& archive,
                          std::vector<std::string>& gathered_archive,
                          lbann_comm& comm);
  void all_gather_archive_new(const std::string& archive,
                              std::vector<std::string>& gathered_archive,
                              lbann_comm& comm);

  template <typename T>
  size_t
  all_gather_field(T data, std::vector<T>& gathered_data, lbann_comm& comm);
  virtual void all_gather_packed_lists(lbann_comm& comm);

  void keep_sample_order(bool keep);

  void set_sample_list_name(const std::string& n);

  void set_data_file_check();
  void unset_data_file_check();

  void build_sample_map_from_name_to_index();

  void clear_sample_map_from_name_to_index();

  sample_idx_t get_sample_index(const sample_name_t& sn);

protected:
  std::string read_header_line(std::istream& ifs,
                               const std::string& listname,
                               const std::string& info);

  void read_header(std::istream& istrm);

  virtual void
  read_sample_list(std::istream& istrm, size_t stride = 1, size_t offset = 0);

  virtual void assign_samples_name();

  size_t get_samples_per_file(std::istream& istrm,
                              size_t stride = 1,
                              size_t offset = 0);

  void write_header(std::string& sstr, size_t num_files) const;

  virtual void
  get_num_samples(size_t& total, size_t& included, size_t& excluded) const;

  virtual void set_samples_filename(sample_file_id_t id,
                                    const std::string& filename);

  virtual void reorder();

protected:
  sample_list_header m_header;

  size_t m_stride;

  bool m_keep_order;

  bool m_check_data_file;

  samples_t m_sample_list;

  sample_map_t m_map_name_to_idx;

private:
  file_id_stats_v_t m_file_id_stats_map;
};

void handle_mpi_error(int ierr);

template <typename T>
inline T uninitialized_sample_name();

} // namespace lbann

#endif // LBANN_DATA_READERS_SAMPLE_LIST_HPP