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             : #include "LibtorchArtificialNeuralNet.h"
      13             : #include "MooseError.h"
      14             : 
      15             : namespace Moose
      16             : {
      17             : 
      18          58 : LibtorchArtificialNeuralNet::LibtorchArtificialNeuralNet(
      19             :     const std::string name,
      20             :     const unsigned int num_inputs,
      21             :     const unsigned int num_outputs,
      22             :     const std::vector<unsigned int> & num_neurons_per_layer,
      23             :     const std::vector<std::string> & activation_function,
      24             :     const torch::DeviceType device_type,
      25          58 :     const torch::ScalarType data_type)
      26          58 :   : _name(name),
      27          58 :     _num_inputs(num_inputs),
      28          58 :     _num_outputs(num_outputs),
      29          58 :     _num_neurons_per_layer(num_neurons_per_layer),
      30          58 :     _activation_function(MultiMooseEnum("relu sigmoid elu gelu linear", "relu")),
      31          58 :     _device_type(device_type),
      32          58 :     _data_type(data_type)
      33             : {
      34          58 :   _activation_function = activation_function;
      35             : 
      36             :   // Check if the number of activation functions matches the number of hidden layers
      37          88 :   if ((_activation_function.size() != 1) &&
      38          31 :       (_activation_function.size() != _num_neurons_per_layer.size()))
      39           1 :     mooseError("The number of activation functions should be either one or the same as the number "
      40             :                "of hidden layers");
      41          56 :   constructNeuralNetwork();
      42          56 : }
      43             : 
      44           1 : LibtorchArtificialNeuralNet::LibtorchArtificialNeuralNet(
      45           1 :     const Moose::LibtorchArtificialNeuralNet & nn)
      46             :   : torch::nn::Module(),
      47           1 :     _name(nn.name()),
      48           1 :     _num_inputs(nn.numInputs()),
      49           1 :     _num_outputs(nn.numOutputs()),
      50           1 :     _num_neurons_per_layer(nn.numNeuronsPerLayer()),
      51           1 :     _activation_function(nn.activationFunctions()),
      52           1 :     _device_type(nn.deviceType()),
      53           3 :     _data_type(nn.dataType())
      54             : {
      55             : 
      56             :   // We construct the NN architecture
      57           1 :   constructNeuralNetwork();
      58             :   // We fill it up with the current parameter values
      59           1 :   const auto & from_params = nn.named_parameters();
      60           1 :   auto to_params = this->named_parameters();
      61           7 :   for (unsigned int param_i : make_range(from_params.size()))
      62           6 :     to_params[param_i].value().data() = from_params[param_i].value().data().clone();
      63           1 : }
      64             : 
      65             : void
      66          57 : LibtorchArtificialNeuralNet::constructNeuralNetwork()
      67             : {
      68             :   // Adding hidden layers
      69          57 :   unsigned int inp_neurons = _num_inputs;
      70         196 :   for (unsigned int i = 0; i < numHiddenLayers(); ++i)
      71             :   {
      72             :     std::unordered_map<std::string, unsigned int> parameters = {
      73         417 :         {"inp_neurons", inp_neurons}, {"out_neurons", _num_neurons_per_layer[i]}};
      74         139 :     addLayer("hidden_layer_" + std::to_string(i + 1), parameters);
      75             : 
      76             :     // Necessary to retain double precision (and error-free runs)
      77         139 :     _weights[i]->to(_device_type, _data_type);
      78         139 :     inp_neurons = _num_neurons_per_layer[i];
      79         139 :   }
      80             :   // Adding output layer
      81             :   std::unordered_map<std::string, unsigned int> parameters = {{"inp_neurons", inp_neurons},
      82         171 :                                                               {"out_neurons", _num_outputs}};
      83          57 :   addLayer("output_layer_", parameters);
      84          57 :   _weights.back()->to(_device_type, _data_type);
      85         253 : }
      86             : 
      87             : torch::Tensor
      88       40909 : LibtorchArtificialNeuralNet::forward(const torch::Tensor & x)
      89             : {
      90       40909 :   torch::Tensor output(x);
      91       40909 :   if (_data_type != output.scalar_type())
      92           0 :     output.to(_data_type);
      93       40909 :   if (_device_type != output.device().type())
      94           0 :     output.to(_device_type);
      95             : 
      96      163551 :   for (unsigned int i = 0; i < _weights.size() - 1; ++i)
      97             :   {
      98             :     std::string activation =
      99      122642 :         _activation_function.size() > 1 ? _activation_function[i] : _activation_function[0];
     100      122642 :     if (activation == "relu")
     101      122546 :       output = torch::relu(_weights[i]->forward(output));
     102          96 :     else if (activation == "sigmoid")
     103          24 :       output = torch::sigmoid(_weights[i]->forward(output));
     104          72 :     else if (activation == "elu")
     105          24 :       output = torch::elu(_weights[i]->forward(output));
     106          48 :     else if (activation == "gelu")
     107          24 :       output = torch::gelu(_weights[i]->forward(output));
     108          24 :     else if (activation == "linear")
     109          24 :       output = _weights[i]->forward(output);
     110      122642 :   }
     111             : 
     112       40909 :   output = _weights[_weights.size() - 1]->forward(output);
     113             : 
     114       40909 :   return output;
     115           0 : }
     116             : 
     117             : void
     118         196 : LibtorchArtificialNeuralNet::addLayer(
     119             :     const std::string & layer_name,
     120             :     const std::unordered_map<std::string, unsigned int> & parameters)
     121             : {
     122         196 :   auto it = parameters.find("inp_neurons");
     123         196 :   if (it == parameters.end())
     124           0 :     ::mooseError("Number of input neurons not found during the construction of "
     125             :                  "LibtorchArtificialNeuralNet!");
     126         196 :   unsigned int inp_neurons = it->second;
     127             : 
     128         196 :   it = parameters.find("out_neurons");
     129         196 :   if (it == parameters.end())
     130           0 :     ::mooseError("Number of output neurons not found during the construction of "
     131             :                  "LibtorchArtificialNeuralNet!");
     132         196 :   unsigned int out_neurons = it->second;
     133             : 
     134         196 :   _weights.push_back(register_module(layer_name, torch::nn::Linear(inp_neurons, out_neurons)));
     135         196 : }
     136             : 
     137             : void
     138           1 : LibtorchArtificialNeuralNet::store(nlohmann::json & json) const
     139             : {
     140           1 :   const auto & named_params = this->named_parameters();
     141           7 :   for (const auto & param_i : make_range(named_params.size()))
     142             :   {
     143             :     // We cast the parameters into a 1D vector
     144          30 :     json[named_params[param_i].key()] = std::vector<Real>(
     145           6 :         named_params[param_i].value().data_ptr<Real>(),
     146          30 :         named_params[param_i].value().data_ptr<Real>() + named_params[param_i].value().numel());
     147             :   }
     148           1 : }
     149             : 
     150             : void
     151           1 : to_json(nlohmann::json & json, const Moose::LibtorchArtificialNeuralNet * const & network)
     152             : {
     153           1 :   if (network)
     154           1 :     network->store(json);
     155           1 : }
     156             : 
     157             : }
     158             : 
     159             : template <>
     160             : void
     161           0 : dataStore<Moose::LibtorchArtificialNeuralNet>(
     162             :     std::ostream & stream, std::shared_ptr<Moose::LibtorchArtificialNeuralNet> & nn, void * context)
     163             : {
     164           0 :   std::string n(nn->name());
     165           0 :   dataStore(stream, n, context);
     166             : 
     167           0 :   unsigned int ni(nn->numInputs());
     168           0 :   dataStore(stream, ni, context);
     169             : 
     170           0 :   unsigned int no(nn->numOutputs());
     171           0 :   dataStore(stream, no, context);
     172             : 
     173           0 :   unsigned int nhl(nn->numHiddenLayers());
     174           0 :   dataStore(stream, nhl, context);
     175             : 
     176           0 :   std::vector<unsigned int> nnpl(nn->numNeuronsPerLayer());
     177           0 :   dataStore(stream, nnpl, context);
     178             : 
     179           0 :   unsigned int afs(nn->activationFunctions().size());
     180           0 :   dataStore(stream, afs, context);
     181             : 
     182           0 :   std::vector<std::string> items(afs);
     183           0 :   for (unsigned int i = 0; i < afs; ++i)
     184           0 :     items[i] = nn->activationFunctions()[i];
     185             : 
     186           0 :   dataStore(stream, items, context);
     187             : 
     188           0 :   auto device_type = static_cast<std::underlying_type<torch::DeviceType>::type>(nn->deviceType());
     189           0 :   dataStore(stream, device_type, context);
     190             : 
     191           0 :   auto data_type = static_cast<std::underlying_type<torch::ScalarType>::type>(nn->dataType());
     192           0 :   dataStore(stream, data_type, context);
     193             : 
     194           0 :   torch::save(nn, nn->name());
     195           0 : }
     196             : 
     197             : template <>
     198             : void
     199           0 : dataLoad<Moose::LibtorchArtificialNeuralNet>(
     200             :     std::istream & stream, std::shared_ptr<Moose::LibtorchArtificialNeuralNet> & nn, void * context)
     201             : {
     202           0 :   std::string name;
     203           0 :   dataLoad(stream, name, context);
     204             : 
     205             :   unsigned int num_inputs;
     206           0 :   dataLoad(stream, num_inputs, context);
     207             : 
     208             :   unsigned int num_outputs;
     209           0 :   dataLoad(stream, num_outputs, context);
     210             : 
     211             :   unsigned int num_hidden_layers;
     212           0 :   dataLoad(stream, num_hidden_layers, context);
     213             : 
     214           0 :   std::vector<unsigned int> num_neurons_per_layer;
     215           0 :   num_neurons_per_layer.resize(num_hidden_layers);
     216           0 :   dataLoad(stream, num_neurons_per_layer, context);
     217             : 
     218             :   unsigned int num_activation_items;
     219           0 :   dataLoad(stream, num_activation_items, context);
     220             : 
     221           0 :   std::vector<std::string> activation_functions;
     222           0 :   activation_functions.resize(num_activation_items);
     223           0 :   dataLoad(stream, activation_functions, context);
     224             : 
     225             :   std::underlying_type<torch::DeviceType>::type device_type;
     226           0 :   dataLoad(stream, device_type, context);
     227           0 :   const torch::DeviceType divt(static_cast<torch::DeviceType>(device_type));
     228             : 
     229             :   std::underlying_type<torch::ScalarType>::type data_type;
     230           0 :   dataLoad(stream, data_type, context);
     231           0 :   const torch::ScalarType datt(static_cast<torch::ScalarType>(data_type));
     232             : 
     233           0 :   nn = std::make_shared<Moose::LibtorchArtificialNeuralNet>(
     234           0 :       name, num_inputs, num_outputs, num_neurons_per_layer, activation_functions, divt, datt);
     235             : 
     236           0 :   torch::load(nn, name);
     237           0 : }
     238             : 
     239             : template <>
     240             : void
     241           0 : dataStore<Moose::LibtorchArtificialNeuralNet const>(
     242             :     std::ostream & /*stream*/,
     243             :     Moose::LibtorchArtificialNeuralNet const *& /*nn*/,
     244             :     void * /*context*/)
     245             : {
     246           0 : }
     247             : 
     248             : template <>
     249             : void
     250           0 : dataLoad<Moose::LibtorchArtificialNeuralNet const>(
     251             :     std::istream & /*stream*/,
     252             :     Moose::LibtorchArtificialNeuralNet const *& /*nn*/,
     253             :     void * /*context*/)
     254             : {
     255           0 : }
     256             : 
     257             : #endif