Line data    Source code

       1             : //* This file is part of the MOOSE framework
       2             : //* https://mooseframework.inl.gov
       3             : //*
       4             : //* All rights reserved, see COPYRIGHT for full restrictions
       5             : //* https://github.com/idaholab/moose/blob/master/COPYRIGHT
       6             : //*
       7             : //* Licensed under LGPL 2.1, please see LICENSE for details
       8             : //* https://www.gnu.org/licenses/lgpl-2.1.html
       9             : 
      10             : #ifdef LIBTORCH_ENABLED
      11             : 
      12             : #pragma once
      13             : 
      14             : #include <torch/torch.h>
      15             : #include <torch/script.h>
      16             : #include "LibtorchNeuralNetBase.h"
      17             : #include "MooseError.h"
      18             : #include "DataIO.h"
      19             : #include "MultiMooseEnum.h"
      20             : #include "nlohmann/json.h"
      21             : 
      22             : namespace Moose
      23             : {
      24             : 
      25             : // A class that describes a simple feed-forward neural net.
      26             : class LibtorchArtificialNeuralNet : public torch::nn::Module, public LibtorchNeuralNetBase
      27             : {
      28             : public:
      29             :   /**
      30             :    * Construct using input parameters
      31             :    * @param name Name of the neural network
      32             :    * @param num_inputs The number of input neurons/parameters
      33             :    * @param num_neurons_per_layer Number of neurons per hidden layer
      34             :    * @param num_outputs The number of output neurons
      35             :    */
      36             :   LibtorchArtificialNeuralNet(const std::string name,
      37             :                               const unsigned int num_inputs,
      38             :                               const unsigned int num_outputs,
      39             :                               const std::vector<unsigned int> & num_neurons_per_layer,
      40             :                               const std::vector<std::string> & activation_function = {"relu"},
      41             :                               const torch::DeviceType device_type = torch::kCPU,
      42             :                               const torch::ScalarType scalar_type = torch::kDouble);
      43             : 
      44             :   /**
      45             :    * Copy construct an artificial neural network
      46             :    * @param nn The neural network which needs to be copied
      47             :    */
      48             :   LibtorchArtificialNeuralNet(const Moose::LibtorchArtificialNeuralNet & nn);
      49             : 
      50             :   /**
      51             :    * Add layers to the neural network
      52             :    * @param layer_name The name of the layer to be added
      53             :    * @param parameters A map of parameter names and the corresponding values which
      54             :    *                   describe the neural net layer architecture
      55             :    */
      56             :   virtual void addLayer(const std::string & layer_name,
      57             :                         const std::unordered_map<std::string, unsigned int> & parameters);
      58             : 
      59             :   /**
      60             :    * Overriding the forward substitution function for the neural network, unfortunately
      61             :    * this cannot be const since it creates a graph in the background
      62             :    * @param x Input tensor for the evaluation
      63             :    */
      64             :   virtual torch::Tensor forward(const torch::Tensor & x) override;
      65             : 
      66             :   /// Return the name of the neural network
      67           1 :   const std::string & name() const { return _name; }
      68             :   /// Return the number of neurons on the input layer
      69           1 :   unsigned int numInputs() const { return _num_inputs; }
      70             :   /// Return the number of neurons on the output layer
      71           1 :   unsigned int numOutputs() const { return _num_outputs; }
      72             :   /// Return the number of hidden layers
      73         196 :   unsigned int numHiddenLayers() const { return _num_neurons_per_layer.size(); }
      74             :   /// Return the hidden layer architecture
      75           1 :   const std::vector<unsigned int> & numNeuronsPerLayer() const { return _num_neurons_per_layer; }
      76             :   /// Return the multi enum containing the activation functions
      77           1 :   const MultiMooseEnum & activationFunctions() const { return _activation_function; }
      78             :   /// Return the device which is used by this neural network
      79           1 :   torch::DeviceType deviceType() const { return _device_type; }
      80             :   /// Return the data type which is used by this neural network
      81           1 :   torch::ScalarType dataType() const { return _data_type; }
      82             :   /// Construct the neural network
      83             :   void constructNeuralNetwork();
      84             : 
      85             :   /// Store the network architecture in a json file (for debugging, visualization)
      86             :   void store(nlohmann::json & json) const;
      87             : 
      88             : protected:
      89             :   /// Name of the neural network
      90             :   const std::string _name;
      91             :   /// Submodules that hold linear operations and the corresponding
      92             :   /// weights and biases (y = W * x + b)
      93             :   std::vector<torch::nn::Linear> _weights;
      94             :   // Number of neurons on the input layer
      95             :   const unsigned int _num_inputs;
      96             :   /// Number of neurons on the output layer
      97             :   const unsigned int _num_outputs;
      98             :   /// Hidden layer architecture
      99             :   const std::vector<unsigned int> _num_neurons_per_layer;
     100             :   /// Activation functions (either one for all hidden layers or one for every layer
     101             :   /// separately)
     102             :   MultiMooseEnum _activation_function;
     103             :   /// The device type used for this neural network
     104             :   const torch::DeviceType _device_type;
     105             :   /// The data type used in this neural network
     106             :   const torch::ScalarType _data_type;
     107             : };
     108             : 
     109             : void to_json(nlohmann::json & json, const Moose::LibtorchArtificialNeuralNet * const & network);
     110             : 
     111             : }
     112             : 
     113             : template <>
     114             : void dataStore<Moose::LibtorchArtificialNeuralNet>(
     115             :     std::ostream & stream,
     116             :     std::shared_ptr<Moose::LibtorchArtificialNeuralNet> & nn,
     117             :     void * context);
     118             : 
     119             : template <>
     120             : void dataLoad<Moose::LibtorchArtificialNeuralNet>(
     121             :     std::istream & stream,
     122             :     std::shared_ptr<Moose::LibtorchArtificialNeuralNet> & nn,
     123             :     void * context);
     124             : 
     125             : // This is needed because the reporter which is used to ouput the neural net parameters to JSON
     126             : // requires a dataStore/dataLoad. However, these functions will be empty due to the fact that
     127             : // we are only interested in the JSON output and we don't want to output everything
     128             : template <>
     129             : void dataStore<Moose::LibtorchArtificialNeuralNet const>(
     130             :     std::ostream & stream, Moose::LibtorchArtificialNeuralNet const *& nn, void * context);
     131             : 
     132             : template <>
     133             : void dataLoad<Moose::LibtorchArtificialNeuralNet const>(
     134             :     std::istream & stream, Moose::LibtorchArtificialNeuralNet const *& nn, void * context);
     135             : 
     136             : #endif