data_reader_csv.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.
//
// data_reader_csv .hpp .cpp - generic_data_reader class for CSV files

#ifndef LBANN_DATA_READER_CSV_HPP
#define LBANN_DATA_READER_CSV_HPP

#include "data_reader.hpp"
#include <unordered_map>

namespace lbann {

class csv_reader : public generic_data_reader
{
public:
  csv_reader(bool shuffle = true);
  csv_reader(const csv_reader&);
  csv_reader& operator=(const csv_reader&);
  ~csv_reader() override;

  csv_reader* copy() const override { return new csv_reader(*this); }

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

  void set_label_col(int col) { m_label_col = col; }
  void set_response_col(int col) { m_response_col = col; }
  void disable_labels(bool b = true)
  {
    m_supported_input_types[INPUT_DATA_TYPE_LABELS] = false;
    m_disable_labels = b;
  }
  void enable_responses(bool b = false)
  {
    m_supported_input_types[INPUT_DATA_TYPE_RESPONSES] = true;
    m_disable_responses = b;
  }
  void set_separator(char sep) { m_separator = sep; }
  void set_skip_cols(int cols) { m_skip_cols = cols; }
  void set_skip_rows(int rows) { m_skip_rows = rows; }
  void set_has_header(bool b) { m_has_header = b; }

  void set_column_transform(int col,
                            std::function<DataType(const std::string&)> f)
  {
    m_col_transforms[col] = f;
  }

  void set_label_transform(std::function<int(const std::string&)> f)
  {
    m_label_transform = f;
  }
  void set_response_transform(std::function<DataType(const std::string&)> f)
  {
    m_response_transform = f;
  }

  void load() override;

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

  int get_num_labels() const override { return m_num_labels; }
  int get_linearized_data_size() const override
  {
    // Account for label and skipped columns.
    if (m_label_col < m_skip_cols) {
      return m_num_cols - m_skip_cols;
    }
    else {
      return m_num_cols - 1 - m_skip_cols;
    }
  }
  int get_linearized_label_size() const override { return m_num_labels; }
  const std::vector<El::Int> get_data_dims() const override
  {
    return {get_linearized_data_size()};
  }

  std::vector<DataType> fetch_line_label_response(int data_id);

protected:
  bool fetch_datum(CPUMat& X, int data_id, int mb_idx) override;
  bool fetch_label(CPUMat& Y, int data_id, int mb_idx) override;
  bool fetch_response(CPUMat& Y, int data_id, int mb_idx) override;

  std::vector<DataType> fetch_line(int data_id);

  void skip_rows(std::ifstream& s, int rows);

  void setup_ifstreams();

  std::string fetch_raw_line(int data_id);

  char m_separator = ',';
  int m_skip_cols = 0;
  int m_skip_rows = 0;
  bool m_has_header = true;
  int m_label_col = -1;
  int m_response_col = -1;
  bool m_disable_labels = false;
  bool m_disable_responses = true;
  int m_num_cols = 0;
  int m_num_samples = 0;
  int m_num_labels = 0;
  std::vector<std::ifstream*> m_ifstreams;
  std::vector<std::streampos> m_index;
  std::vector<int> m_labels;
  std::vector<DataType> m_responses;
  std::unordered_map<int, std::function<DataType(const std::string&)>>
    m_col_transforms;
  std::function<int(const std::string&)> m_label_transform =
    [](const std::string& s) -> int { return std::stoi(s); };
  std::function<DataType(const std::string&)> m_response_transform =
    [](const std::string& s) -> DataType { return std::stod(s); };
};

} // namespace lbann

#endif // LBANN_DATA_READER_CSV_HPP