LibtorchArtificialNeuralNetTrainer.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://mooseframework.inl.gov
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html

#ifdef LIBTORCH_ENABLED

#pragma once

#include "LibtorchArtificialNeuralNet.h"
#include "LibtorchDataset.h"
#include "DataIO.h"
#include "MooseEnum.h"

namespace Moose
{

struct LibtorchTrainingOptions
{
  MooseEnum optimizer_type = MooseEnum("adam=0 adagrad=1 rmsprop=2 sgd=3", "adam");
  unsigned int num_epochs = 1;
  unsigned int num_batches = 1;
  bool print_loss = false;
  unsigned int print_epoch_loss = 1;
  Real rel_loss_tol = 1e-12;
  Real learning_rate = 1e-3;
  unsigned int parallel_processes = 1;
  bool allow_duplicates = false;
};

template <typename SamplerType = torch::data::samplers::DistributedSequentialSampler>
class LibtorchArtificialNeuralNetTrainer : public libMesh::ParallelObject
{
public:
  LibtorchArtificialNeuralNetTrainer(LibtorchArtificialNeuralNet & nn,
                                     const Parallel::Communicator & comm);

  virtual void train(LibtorchDataset & dataset, const LibtorchTrainingOptions & options);

  static unsigned int computeBatchSize(const unsigned int num_samples,
                                       const unsigned int num_batches);
  static unsigned int computeLocalBatchSize(const unsigned int batch_size,
                                            const unsigned int num_ranks);

  static std::unique_ptr<torch::optim::Optimizer>
  createOptimizer(const LibtorchArtificialNeuralNet & nn, const LibtorchTrainingOptions & options);

protected:
  LibtorchArtificialNeuralNet & _nn;
};
} // end Moose namespace

#endif