CSVLogger.hpp

#pragma once
 
#include <fstream>
#include <stdexcept>
#include <string>
#include <vector>
 
#include "Callback.hpp"
#include "utils/Functions.hpp"  // fileExistsWithExtension
 
namespace NeuralNet {
class CSVLogger : public Callback {
 public:
  CSVLogger(const std::string &filepath, const std::string &separator = ",") {
    assert(fileHasExtension(filepath, ".csv") &&
           "filepath must have .csv extension");
    this->filepath = filepath;
    this->separator = separator;
  };
 
  void onEpochBegin(Model &model) override {};
 
  void onEpochEnd(Model &model) override {
    std::unordered_map<std::string, Logs> logs = getLogs(model);
    std::vector<double> row;
 
    row.reserve(logs.size());
 
    std::transform(logs.begin(), logs.end(), std::back_inserter(row),
                   [](const auto &log) {
                     const auto &value = log.second;
                     if (std::holds_alternative<int>(value)) {
                       return static_cast<double>(std::get<int>(value));
                     }
                     return std::get<double>(value);
                   });
 
    data.push_back(row);
  };
 
  void onTrainBegin(Model &model) override {
    std::unordered_map<std::string, Logs> logs = getLogs(model);
    // Initializing the headers with the logs keys
    for (const auto &log : logs) {
      const auto &value = log.second;
      if (std::holds_alternative<int>(value) ||
          std::holds_alternative<double>(value)) {
        headers.push_back(log.first);
      }
    };
  };
 
  void onTrainEnd(Model &model) override {
    std::ofstream file(filepath);
 
    if (!file.is_open()) {
      throw std::runtime_error("Couldn't open csv file");
    }
 
    file << formatRow(headers);
 
    std::for_each(data.begin(), data.end(),
                  [&file, this](auto &row) { file << this->formatRow(row); });
 
    file.close();
  };
 
  void onBatchBegin(Model &model) override {};
  void onBatchEnd(Model &model) override {};
 
 private:
  std::string filepath;
  std::string separator;
  std::vector<std::string> headers;
  std::vector<std::vector<double>> data;
 
  template <typename T>
  std::string formatRow(const std::vector<T> &v) {
    std::string csvRow;
 
    for (T el : v) {
      csvRow += std::to_string(el) + separator;
    }
 
    // Remove last ","
    csvRow.pop_back();
 
    return csvRow + "\n";
  };
 
  std::string formatRow(const std::vector<std::string> &v) {
    std::string csvRow;
 
    for (const std::string &el : v) {
      csvRow += el + separator;
    }
 
    // Remove last ","
    csvRow.pop_back();
 
    return csvRow + "\n";
  };
};
}  // namespace NeuralNet