#include <LibtorchArtificialNeuralNet.h>

Inheritance diagram for Moose::LibtorchArtificialNeuralNet:

Public Member Functions
	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)
	Construct using input parameters.

	LibtorchArtificialNeuralNet (const Moose::LibtorchArtificialNeuralNet &nn)
	Copy construct an artificial neural network.

virtual void	addLayer (const std::string &layer_name, const std::unordered_map< std::string, unsigned int > &parameters)
	Add layers to the neural network.

virtual torch::Tensor	forward (const torch::Tensor &x) override
	Overriding the forward substitution function for the neural network, unfortunately this cannot be const since it creates a graph in the background.

const std::string &	name () const
	Return the name of the neural network.

unsigned int	numInputs () const
	Return the number of neurons on the input layer.

unsigned int	numOutputs () const
	Return the number of neurons on the output layer.

unsigned int	numHiddenLayers () const
	Return the number of hidden layers.

const std::vector< unsigned int > &	numNeuronsPerLayer () const
	Return the hidden layer architecture.

const MultiMooseEnum &	activationFunctions () const
	Return the multi enum containing the activation functions.

torch::DeviceType	deviceType () const
	Return the device which is used by this neural network.

torch::ScalarType	dataType () const
	Return the data type which is used by this neural network.

void	constructNeuralNetwork ()
	Construct the neural network.

void	store (nlohmann::json &json) const
	Store the network architecture in a json file (for debugging, visualization)

Protected Attributes
const std::string	_name
	Name of the neural network.

std::vector< torch::nn::Linear >	_weights
	Submodules that hold linear operations and the corresponding weights and biases (y = W * x + b)

const unsigned int	_num_inputs

const unsigned int	_num_outputs
	Number of neurons on the output layer.

const std::vector< unsigned int >	_num_neurons_per_layer
	Hidden layer architecture.

MultiMooseEnum	_activation_function
	Activation functions (either one for all hidden layers or one for every layer separately)

const torch::DeviceType	_device_type
	The device type used for this neural network.

const torch::ScalarType	_data_type
	The data type used in this neural network.

Detailed Description

Definition at line 26 of file LibtorchArtificialNeuralNet.h.

Constructor & Destructor Documentation

◆ LibtorchArtificialNeuralNet() [1/2]

Moose::LibtorchArtificialNeuralNet::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`
	)

Construct using input parameters.

Parameters

name	Name of the neural network
num_inputs	The number of input neurons/parameters
num_neurons_per_layer	Number of neurons per hidden layer
num_outputs	The number of output neurons

Definition at line 18 of file LibtorchArtificialNeuralNet.C.

  : _name(name),
    _num_inputs(num_inputs),
    _num_outputs(num_outputs),
    _num_neurons_per_layer(num_neurons_per_layer),
    _activation_function(MultiMooseEnum("relu sigmoid elu gelu linear", "relu")),
    _device_type(device_type),
    _data_type(data_type)
{
  _activation_function = activation_function;
 
  // Check if the number of activation functions matches the number of hidden layers
  if ((_activation_function.size() != 1) &&
      (_activation_function.size() != _num_neurons_per_layer.size()))
    mooseError("The number of activation functions should be either one or the same as the number "
               "of hidden layers");
  constructNeuralNetwork();
}

◆ LibtorchArtificialNeuralNet() [2/2]

Moose::LibtorchArtificialNeuralNet::LibtorchArtificialNeuralNet ( const Moose::LibtorchArtificialNeuralNet & nn )

Copy construct an artificial neural network.

Parameters

nn	The neural network which needs to be copied

Definition at line 44 of file LibtorchArtificialNeuralNet.C.

  : torch::nn::Module(),
    _name(nn.name()),
    _num_inputs(nn.numInputs()),
    _num_outputs(nn.numOutputs()),
    _num_neurons_per_layer(nn.numNeuronsPerLayer()),
    _activation_function(nn.activationFunctions()),
    _device_type(nn.deviceType()),
    _data_type(nn.dataType())
{
 
  // We construct the NN architecture
  constructNeuralNetwork();
  // We fill it up with the current parameter values
  const auto & from_params = nn.named_parameters();
  auto to_params = this->named_parameters();
  for (unsigned int param_i : make_range(from_params.size()))
    to_params[param_i].value().data() = from_params[param_i].value().data().clone();
}

Member Function Documentation

◆ activationFunctions()

const MultiMooseEnum & Moose::LibtorchArtificialNeuralNet::activationFunctions ( ) const

inline

Return the multi enum containing the activation functions.

Definition at line 77 of file LibtorchArtificialNeuralNet.h.

77{ return _activation_function; }

◆ addLayer()

void Moose::LibtorchArtificialNeuralNet::addLayer	(	const std::string &	layer_name,
		const std::unordered_map< std::string, unsigned int > &	parameters
	)

virtual

Add layers to the neural network.

Parameters

layer_name	The name of the layer to be added
parameters	A map of parameter names and the corresponding values which describe the neural net layer architecture

Definition at line 118 of file LibtorchArtificialNeuralNet.C.

{
  auto it = parameters.find("inp_neurons");
  if (it == parameters.end())
    ::mooseError("Number of input neurons not found during the construction of "
                 "LibtorchArtificialNeuralNet!");
  unsigned int inp_neurons = it->second;
 
  it = parameters.find("out_neurons");
  if (it == parameters.end())
    ::mooseError("Number of output neurons not found during the construction of "
                 "LibtorchArtificialNeuralNet!");
  unsigned int out_neurons = it->second;
 
  _weights.push_back(register_module(layer_name, torch::nn::Linear(inp_neurons, out_neurons)));
}

Referenced by constructNeuralNetwork().

◆ constructNeuralNetwork()

void Moose::LibtorchArtificialNeuralNet::constructNeuralNetwork ( )

Construct the neural network.

Definition at line 66 of file LibtorchArtificialNeuralNet.C.

{
  // Adding hidden layers
  unsigned int inp_neurons = _num_inputs;
  for (unsigned int i = 0; i < numHiddenLayers(); ++i)
  {
    std::unordered_map<std::string, unsigned int> parameters = {
        {"inp_neurons", inp_neurons}, {"out_neurons", _num_neurons_per_layer[i]}};
    addLayer("hidden_layer_" + std::to_string(i + 1), parameters);
 
    // Necessary to retain double precision (and error-free runs)
    _weights[i]->to(_device_type, _data_type);
    inp_neurons = _num_neurons_per_layer[i];
  }
  // Adding output layer
  std::unordered_map<std::string, unsigned int> parameters = {{"inp_neurons", inp_neurons},
                                                              {"out_neurons", _num_outputs}};
  addLayer("output_layer_", parameters);
  _weights.back()->to(_device_type, _data_type);
}

Referenced by LibtorchArtificialNeuralNet(), and LibtorchArtificialNeuralNet().

◆ dataType()

torch::ScalarType Moose::LibtorchArtificialNeuralNet::dataType ( ) const

inline

Return the data type which is used by this neural network.

Definition at line 81 of file LibtorchArtificialNeuralNet.h.

81{ return _data_type; }

◆ deviceType()

torch::DeviceType Moose::LibtorchArtificialNeuralNet::deviceType ( ) const

inline

Return the device which is used by this neural network.

Definition at line 79 of file LibtorchArtificialNeuralNet.h.

79{ return _device_type; }

◆ forward()

torch::Tensor Moose::LibtorchArtificialNeuralNet::forward ( const torch::Tensor & x )

overridevirtual

Overriding the forward substitution function for the neural network, unfortunately this cannot be const since it creates a graph in the background.

Parameters

x	Input tensor for the evaluation

Implements Moose::LibtorchNeuralNetBase.

Definition at line 88 of file LibtorchArtificialNeuralNet.C.

{
  torch::Tensor output(x);
  if (_data_type != output.scalar_type())
    output.to(_data_type);
  if (_device_type != output.device().type())
    output.to(_device_type);
 
  for (unsigned int i = 0; i < _weights.size() - 1; ++i)
  {
    std::string activation =
        _activation_function.size() > 1 ? _activation_function[i] : _activation_function[0];
    if (activation == "relu")
      output = torch::relu(_weights[i]->forward(output));
    else if (activation == "sigmoid")
      output = torch::sigmoid(_weights[i]->forward(output));
    else if (activation == "elu")
      output = torch::elu(_weights[i]->forward(output));
    else if (activation == "gelu")
      output = torch::gelu(_weights[i]->forward(output));
    else if (activation == "linear")
      output = _weights[i]->forward(output);
  }
 
  output = _weights[_weights.size() - 1]->forward(output);
 
  return output;
}

Referenced by forward().

◆ name()

const std::string & Moose::LibtorchArtificialNeuralNet::name ( ) const

inline

Return the name of the neural network.

Definition at line 67 of file LibtorchArtificialNeuralNet.h.

67{ return _name; }

◆ numHiddenLayers()

unsigned int Moose::LibtorchArtificialNeuralNet::numHiddenLayers ( ) const

inline

Return the number of hidden layers.

Definition at line 73 of file LibtorchArtificialNeuralNet.h.

73{ return _num_neurons_per_layer.size(); }

Referenced by constructNeuralNetwork().

◆ numInputs()

unsigned int Moose::LibtorchArtificialNeuralNet::numInputs ( ) const

inline

Return the number of neurons on the input layer.

Definition at line 69 of file LibtorchArtificialNeuralNet.h.

69{ return _num_inputs; }

◆ numNeuronsPerLayer()

const std::vector< unsigned int > & Moose::LibtorchArtificialNeuralNet::numNeuronsPerLayer ( ) const

inline

Return the hidden layer architecture.

Definition at line 75 of file LibtorchArtificialNeuralNet.h.

75{ return _num_neurons_per_layer; }

◆ numOutputs()

unsigned int Moose::LibtorchArtificialNeuralNet::numOutputs ( ) const

inline

Return the number of neurons on the output layer.

Definition at line 71 of file LibtorchArtificialNeuralNet.h.

71{ return _num_outputs; }

◆ store()

void Moose::LibtorchArtificialNeuralNet::store ( nlohmann::json & json ) const

Store the network architecture in a json file (for debugging, visualization)

Definition at line 138 of file LibtorchArtificialNeuralNet.C.

{
  const auto & named_params = this->named_parameters();
  for (const auto & param_i : make_range(named_params.size()))
  {
    // We cast the parameters into a 1D vector
    json[named_params[param_i].key()] = std::vector<Real>(
        named_params[param_i].value().data_ptr<Real>(),
        named_params[param_i].value().data_ptr<Real>() + named_params[param_i].value().numel());
  }
}