LibtorchArtificialNeuralNet.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 MOOSE_LIBTORCH_ENABLED
 
#pragma once
 
#include <torch/torch.h>
#include <torch/script.h>
#include <torch/serialize/archive.h>
#include "LibtorchNeuralNetBase.h"
#include "MooseError.h"
#include "DataIO.h"
#include "MultiMooseEnum.h"
#include "nlohmann/json.h"
 
namespace Moose
{
 
class LibtorchArtificialNeuralNet : public torch::nn::Module, public LibtorchNeuralNetBase
{
public:
  LibtorchArtificialNeuralNet(const std::string name,
                              const unsigned int num_inputs,
                              const unsigned int num_outputs,
                              const std::vector<unsigned int> & num_neurons_per_layer,
                              const std::vector<std::string> & activation_function = {"relu"},
                              const torch::DeviceType device_type = torch::kCPU,
                              const torch::ScalarType scalar_type = torch::kDouble,
                              const bool build_on_construct = true,
                              const std::vector<Real> & input_shift_factors = {},
                              const std::vector<Real> & input_scaling_factors = {},
                              const std::vector<Real> & output_scaling_factors = {});
 
  LibtorchArtificialNeuralNet(const Moose::LibtorchArtificialNeuralNet & nn,
                              const bool build_on_construct = true);
 
  virtual void addLayer(const std::string & layer_name,
                        const std::unordered_map<std::string, unsigned int> & parameters);
 
  virtual torch::Tensor forward(const torch::Tensor & x) override;
 
  const std::string & name() const { return _name; }
  unsigned int numInputs() const { return _num_inputs; }
  unsigned int numOutputs() const { return _num_outputs; }
  unsigned int numHiddenLayers() const { return _num_neurons_per_layer.size(); }
  const std::vector<unsigned int> & numNeuronsPerLayer() const { return _num_neurons_per_layer; }
  const MultiMooseEnum & activationFunctions() const { return _activation_function; }
  torch::DeviceType deviceType() const { return _device_type; }
  torch::ScalarType dataType() const { return _data_type; }
  const std::vector<Real> & inputShiftFactors() const { return _input_shift_factors; }
  const std::vector<Real> & inputScalingFactors() const { return _input_scaling_factors; }
  const std::vector<Real> & outputScalingFactors() const { return _output_scaling_factors; }
  virtual void constructNeuralNetwork();
 
  void synchronizeAffineFactorsFromBuffers();
 
  Real determineGain(const std::string & activation);
 
  virtual void initializeNeuralNetwork(c10::optional<at::Generator> generator = c10::nullopt);
 
  void store(nlohmann::json & json) const;
 
protected:
  static std::vector<Real> setAffineFactors(const std::vector<Real> & factors,
                                            unsigned int expected_size,
                                            Real default_value,
                                            const std::string & factor_name);
 
  void initializeAffineBuffers();
 
  virtual torch::Tensor preprocessInput(const torch::Tensor & x) const;
 
  virtual torch::Tensor scaleOutput(const torch::Tensor & y) const;
 
  const std::string _name;
  std::vector<torch::nn::Linear> _weights;
  const unsigned int _num_inputs;
  const unsigned int _num_outputs;
  const std::vector<unsigned int> _num_neurons_per_layer;
  MultiMooseEnum _activation_function;
  const torch::DeviceType _device_type;
  const torch::ScalarType _data_type;
  std::vector<Real> _input_shift_factors;
  std::vector<Real> _input_scaling_factors;
  std::vector<Real> _output_scaling_factors;
  torch::Tensor _input_shift_tensor;
  torch::Tensor _input_scale_tensor;
  torch::Tensor _output_scale_tensor;
};
 
void to_json(nlohmann::json & json, const Moose::LibtorchArtificialNeuralNet * const & network);
 
void loadLibtorchArtificialNeuralNetState(Moose::LibtorchArtificialNeuralNet & nn,
                                          const std::string & filename);
}
 
template <>
void dataStore<Moose::LibtorchArtificialNeuralNet>(
    std::ostream & stream,
    std::shared_ptr<Moose::LibtorchArtificialNeuralNet> & nn,
    void * context);
 
template <>
void dataLoad<Moose::LibtorchArtificialNeuralNet>(
    std::istream & stream,
    std::shared_ptr<Moose::LibtorchArtificialNeuralNet> & nn,
    void * context);
 
// This is needed because the reporter which is used to ouput the neural net parameters to JSON
// requires a dataStore/dataLoad. However, these functions will be empty due to the fact that
// we are only interested in the JSON output and we don't want to output everything
template <>
void dataStore<Moose::LibtorchArtificialNeuralNet const>(
    std::ostream & stream, Moose::LibtorchArtificialNeuralNet const *& nn, void * context);
 
template <>
void dataLoad<Moose::LibtorchArtificialNeuralNet const>(
    std::istream & stream, Moose::LibtorchArtificialNeuralNet const *& nn, void * context);
 
#endif