Merge pull request #24893 from idaholab/standardization-for-nn

Add standardization for neural net surrogates.

modules/stochastic_tools/examples/surrogates/cross_validation/all_trainers_uniform_cv.i

@@ -118,15 +118,17 @@
     print_epoch_loss = 0
     activation_function = 'relu'
     cv_surrogate = "ann_surr"
+    standardize_input = false
+    standardize_output = false
   []
 []
 
 [Covariance]
   [rbf]
-    type=SquaredExponentialCovariance
+    type = SquaredExponentialCovariance
     noise_variance = 3.79e-6
-    signal_variance = 1                      #Use a signal variance of 1 in the kernel
-    length_factor = '5.34471 1.41191 5.90721 2.83723'                  #Select a length factor for each parameter
+    signal_variance = 1 #Use a signal variance of 1 in the kernel
+    length_factor = '5.34471 1.41191 5.90721 2.83723' #Select a length factor for each parameter
   []
 []
 

modules/stochastic_tools/include/libtorch/surrogates/LibtorchANNSurrogate.h

@@ -14,6 +14,7 @@
 #include <torch/torch.h>
 #include "LibtorchArtificialNeuralNet.h"
 #include "SurrogateModel.h"
+#include "Standardizer.h"
 
 class LibtorchANNSurrogate : public SurrogateModel
 {
@@ -27,6 +28,12 @@ class LibtorchANNSurrogate : public SurrogateModel
 protected:
   /// Pointer to the neural net object (initialized as null)
   const std::shared_ptr<Moose::LibtorchArtificialNeuralNet> & _nn;
+
+  /// Standardizer for use with input (x)
+  const StochasticTools::Standardizer & _input_standardizer;
+
+  /// Standardizer for use with output response (y)
+  const StochasticTools::Standardizer & _output_standardizer;
 };
 
 #endif

modules/stochastic_tools/include/libtorch/surrogates/LibtorchANNTrainer.h

@@ -16,6 +16,8 @@
 #include "LibtorchArtificialNeuralNetTrainer.h"
 #include "libmesh/utility.h"
 #include "SurrogateTrainer.h"
+#include "Standardizer.h"
+#include "LibtorchUtils.h"
 
 /// Trainer responsible of fitting a neural network on predefined data
 class LibtorchANNTrainer : public SurrogateTrainer
@@ -66,6 +68,18 @@ class LibtorchANNTrainer : public SurrogateTrainer
 
   /// Pointer to the neural net object (initialized as null)
   std::shared_ptr<Moose::LibtorchArtificialNeuralNet> & _nn;
+
+  /// If the training output should be standardized (scaled and shifted)
+  const bool _standardize_input;
+
+  /// If the training output should be standardized (scaled and shifted)
+  const bool _standardize_output;
+
+  /// Standardizer for use with input (x)
+  StochasticTools::Standardizer & _input_standardizer;
+
+  /// Standardizer for use with output response (y)
+  StochasticTools::Standardizer & _output_standardizer;
 };
 
 #endif

modules/stochastic_tools/include/utils/Standardizer.h

@@ -33,6 +33,11 @@ class Standardizer
   /// Sets mean and std directly using provided vectors
   void set(const std::vector<Real> & mean, const std::vector<Real> & stdev);
 
+  /// Get the mean vector
+  const std::vector<Real> & getMean() const { return _mean; }
+  /// Get the standard deviation vector
+  const std::vector<Real> & getStdDev() const { return _stdev; }
+
   /// Methods for computing and setting mean and standard
   void computeSet(const RealEigenMatrix & input);
 

modules/stochastic_tools/src/libtorch/surrogates/LibtorchANNSurrogate.C

@@ -23,7 +23,9 @@ LibtorchANNSurrogate::validParams()
 
 LibtorchANNSurrogate::LibtorchANNSurrogate(const InputParameters & parameters)
   : SurrogateModel(parameters),
-    _nn(getModelData<std::shared_ptr<Moose::LibtorchArtificialNeuralNet>>("nn"))
+    _nn(getModelData<std::shared_ptr<Moose::LibtorchArtificialNeuralNet>>("nn")),
+    _input_standardizer(getModelData<StochasticTools::Standardizer>("input_standardizer")),
+    _output_standardizer(getModelData<StochasticTools::Standardizer>("output_standardizer"))
 {
   // We check if MOOSE is compiled with torch, if not this throws an error
   StochasticToolsApp::requiresTorch(*this);
@@ -38,10 +40,29 @@ LibtorchANNSurrogate::evaluate(const std::vector<Real> & x) const
   mooseAssert(_nn->numInputs() == x.size(),
               "Input point does not match dimensionality of training data.");
 
-  torch::Tensor x_tf = torch::tensor(torch::ArrayRef<Real>(x.data(), x.size())).to(at::kDouble);
+  std::vector<Real> converted_input(x.size(), 0);
+  const auto & input_mean = _input_standardizer.getMean();
+  const auto & input_std = _input_standardizer.getStdDev();
+
+  mooseAssert(mean.size() == converted_input.size() && std.size() == converted_input.size(),
+              "The input standardizer's dimensions should be the same as the input dimension!");
+
+  for (auto input_i : index_range(converted_input))
+    converted_input[input_i] = (x[input_i] - input_mean[input_i]) / input_std[input_i];
+
+  torch::Tensor x_tf =
+      torch::tensor(torch::ArrayRef<Real>(converted_input.data(), converted_input.size()))
+          .to(at::kDouble);
+
+  const auto & output_mean = _output_standardizer.getMean();
+  const auto & output_std = _output_standardizer.getStdDev();
+
+  mooseAssert(output_mean.size() == 1 && output_std.size() == 1,
+              "The output standardizer's dimensions should be 1!");
 
   // Compute prediction
   val = _nn->forward(x_tf).item<double>();
+  val = val * output_std[0] + output_mean[0];
 
   return val;
 }

modules/stochastic_tools/src/libtorch/surrogates/LibtorchANNTrainer.C

@@ -54,6 +54,11 @@ LibtorchANNTrainer::validParams()
   params.addParam<unsigned int>(
       "max_processes", 1, "The maximum number of parallel processes that the trainer will use.");
 
+  params.addParam<bool>(
+      "standardize_input", true, "Standardize (center and scale) training inputs (x values)");
+  params.addParam<bool>(
+      "standardize_output", true, "Standardize (center and scale) training outputs (y values)");
+
   params.suppressParameter<MooseEnum>("response_type");
   return params;
 }
@@ -67,7 +72,11 @@ LibtorchANNTrainer::LibtorchANNTrainer(const InputParameters & parameters)
         "activation_function", getParam<std::vector<std::string>>("activation_function"))),
     _nn_filename(getParam<std::string>("nn_filename")),
     _read_from_file(getParam<bool>("read_from_file")),
-    _nn(declareModelData<std::shared_ptr<Moose::LibtorchArtificialNeuralNet>>("nn"))
+    _nn(declareModelData<std::shared_ptr<Moose::LibtorchArtificialNeuralNet>>("nn")),
+    _standardize_input(getParam<bool>("standardize_input")),
+    _standardize_output(getParam<bool>("standardize_output")),
+    _input_standardizer(declareModelData<StochasticTools::Standardizer>("input_standardizer")),
+    _output_standardizer(declareModelData<StochasticTools::Standardizer>("output_standardizer"))
 {
   // Fixing the RNG seed to make sure every experiment is the same.
   // Otherwise sampling / stochastic gradient descent would be different.
@@ -110,7 +119,7 @@ LibtorchANNTrainer::postTrain()
 
   // Then, we create and load our Tensors
   unsigned int num_samples = _flattened_response.size();
-  unsigned int num_inputs = _sampler.getNumberOfCols();
+  unsigned int num_inputs = _n_dims;
 
   // We create a neural net (for the definition of the net see the header file)
   _nn = std::make_shared<Moose::LibtorchArtificialNeuralNet>(
@@ -135,6 +144,41 @@ LibtorchANNTrainer::postTrain()
   torch::Tensor response_tensor =
       torch::from_blob(_flattened_response.data(), {num_samples, 1}, options).to(at::kDouble);
 
+  // We standardize the input/output pairs if the user requested it
+  if (_standardize_input)
+  {
+    auto data_std_mean = torch::std_mean(data_tensor, 0);
+    auto & data_std = std::get<0>(data_std_mean);
+    auto & data_mean = std::get<1>(data_std_mean);
+
+    data_tensor = (data_tensor - data_mean) / data_std;
+
+    std::vector<Real> converted_data_mean;
+    LibtorchUtils::tensorToVector(data_mean, converted_data_mean);
+    std::vector<Real> converted_data_std;
+    LibtorchUtils::tensorToVector(data_std, converted_data_std);
+    _input_standardizer.set(converted_data_mean, converted_data_std);
+  }
+  else
+    _input_standardizer.set(_n_dims);
+
+  if (_standardize_output)
+  {
+    auto response_std_mean = torch::std_mean(response_tensor, 0);
+    auto & response_std = std::get<0>(response_std_mean);
+    auto & response_mean = std::get<1>(response_std_mean);
+
+    response_tensor = (response_tensor - response_mean) / response_std;
+
+    std::vector<Real> converted_response_mean;
+    LibtorchUtils::tensorToVector(response_mean, converted_response_mean);
+    std::vector<Real> converted_response_std;
+    LibtorchUtils::tensorToVector(response_std, converted_response_std);
+    _output_standardizer.set(converted_response_mean, converted_response_std);
+  }
+  else
+    _output_standardizer.set(1);
+
   // We create a custom data set from our converted data
   Moose::LibtorchDataset my_data(data_tensor, response_tensor);
 

modules/stochastic_tools/test/tests/surrogates/libtorch_nn/gold/evaluate_standardized_results_0002.csv

@@ -0,0 +1,9 @@
+surr
+7.9850031781848
+7.2083697743465
+7.18883960635
+6.4785829501964
+7.1966313094783
+6.4981153459679
+6.4832642706565
+5.8387339623427

modules/stochastic_tools/test/tests/surrogates/libtorch_nn/retrain.i

@@ -38,6 +38,8 @@
     nn_filename = mynet.pt
     read_from_file = false
     print_epoch_loss = 10
+    standardize_input = false
+    standardize_output = false
   []
 []
 

modules/stochastic_tools/test/tests/surrogates/libtorch_nn/tests

@@ -49,4 +49,25 @@
     libtorch = true
     prereq = retrain
   []
+  [train-standardized]
+    requirement = 'The system shall be able to train a libtorch-based neural network with standardized input and output parameters.'
+    type = CheckFiles
+    input = train.i
+    check_files = 'standardized_train.rd mynet.pt'
+    cli_args = "Trainers/train/read_from_file=false Trainers/train/standardize_input=true "
+               "Trainers/train/standardize_output=true Outputs/file_base=standardized"
+    allow_test_objects = true
+    libtorch = true
+  []
+  [evaluate-standardized]
+    requirement = 'The system shall be able to evaluate a previously trained, libtorch-based neural network with standardized input and output parameters.'
+    type = CSVDiff
+    input = evaluate.i
+    allow_test_objects = true
+    csvdiff = 'evaluate_standardized_results_0002.csv'
+    cli_args = "Surrogates/surr/filename='standardized_train.rd' "
+               "Outputs/file_base=evaluate_standardized"
+    prereq = train-standardized
+    libtorch = true
+  []
 []

modules/stochastic_tools/test/tests/surrogates/libtorch_nn/train.i

@@ -34,6 +34,8 @@
     print_epoch_loss = 10
     activation_function = 'relu relu'
     max_processes = 1
+    standardize_input = false
+    standardize_output = false
   []
 []
 

modules/stochastic_tools/test/tests/surrogates/libtorch_nn/train_and_evaluate.i

@@ -39,6 +39,8 @@
     read_from_file = false
     print_epoch_loss = 10
     max_processes = 1
+    standardize_input = false
+    standardize_output = false
   []
 []