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 "LibtorchNeuralNetBase.h"
#include "MooseError.h"
#include "DataIO.h"
#include "MultiMooseEnum.h"
#include "nlohmann/json.h"

namespace Moose
{

// A class that describes a simple feed-forward neural net.
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);

  LibtorchArtificialNeuralNet(const Moose::LibtorchArtificialNeuralNet & nn);

  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; }
  void constructNeuralNetwork();

  void store(nlohmann::json & json) const;

protected:
  const std::string _name;
  std::vector<torch::nn::Linear> _weights;
  // Number of neurons on the input layer
  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;
};

void to_json(nlohmann::json & json, const Moose::LibtorchArtificialNeuralNet * const & network);

}

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