Skip to content
Snippets Groups Projects
Trainer.hpp 481 B
Newer Older
Franck Dary's avatar
Franck Dary committed
#ifndef TRAINER__H
#define TRAINER__H

#include "ReadingMachine.hpp"
#include "ConfigDataset.hpp"
#include "SubConfig.hpp"
#include "TestNetwork.hpp"

class Trainer
{
  private :

  ReadingMachine & machine;
  std::unique_ptr<ConfigDataset> dataset{nullptr};
  std::unique_ptr<torch::optim::Adam> denseOptimizer;
  std::unique_ptr<torch::optim::SparseAdam> sparseOptimizer;

  public :

  Trainer(ReadingMachine & machine);
  void createDataset(SubConfig & goldConfig);
};

#endif