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 : #ifdef MOOSE_LIBTORCH_ENABLED
10 :
11 : #include "GaussianProcessTrainer.h"
12 : #include "Sampler.h"
13 : #include "CartesianProduct.h"
14 :
15 : #include <petsctao.h>
16 : #include <petscdmda.h>
17 :
18 : #include "libmesh/petsc_vector.h"
19 : #include "libmesh/petsc_matrix.h"
20 :
21 : #include <cmath>
22 :
23 : registerMooseObject("StochasticToolsApp", GaussianProcessTrainer);
24 :
25 : InputParameters
26 258 : GaussianProcessTrainer::validParams()
27 : {
28 258 : InputParameters params = SurrogateTrainer::validParams();
29 258 : params.addClassDescription("Provides data preperation and training for a single- or multi-output "
30 : "Gaussian Process surrogate model.");
31 :
32 516 : params.addRequiredParam<UserObjectName>("covariance_function", "Name of covariance function.");
33 516 : params.addParam<bool>(
34 516 : "standardize_params", true, "Standardize (center and scale) training parameters (x values)");
35 516 : params.addParam<bool>(
36 516 : "standardize_data", true, "Standardize (center and scale) training data (y values)");
37 : // Already preparing to use Adam here
38 516 : params.addParam<unsigned int>("num_iters", 1000, "Tolerance value for Adam optimization");
39 516 : params.addParam<unsigned int>("batch_size", 0, "The batch size for Adam optimization");
40 516 : params.addParam<Real>("learning_rate", 0.001, "The learning rate for Adam optimization");
41 516 : params.addParam<MooseEnum>(
42 : "optimizer",
43 774 : MooseEnum("adam=0 legacy_adam=1", "adam"),
44 : "The Adam optimizer semantics to use for Gaussian process hyperparameter tuning.");
45 516 : params.addParam<unsigned int>(
46 : "show_every_nth_iteration",
47 516 : 0,
48 : "Switch to show Adam optimization loss values at every nth step. If 0, nothing is showed.");
49 516 : params.addParam<std::vector<std::string>>("tune_parameters",
50 : "Select hyperparameters to be tuned");
51 516 : params.addParam<std::vector<Real>>("tuning_min", "Minimum allowable tuning value");
52 516 : params.addParam<std::vector<Real>>("tuning_max", "Maximum allowable tuning value");
53 258 : return params;
54 0 : }
55 :
56 130 : GaussianProcessTrainer::GaussianProcessTrainer(const InputParameters & parameters)
57 : : SurrogateTrainer(parameters),
58 : CovarianceInterface(parameters),
59 130 : _predictor_row(getPredictorData()),
60 260 : _gp(declareModelData<StochasticTools::GaussianProcess>("_gp")),
61 260 : _training_params(declareModelData<torch::Tensor>("_training_params")),
62 260 : _standardize_params(getParam<bool>("standardize_params")),
63 260 : _standardize_data(getParam<bool>("standardize_data")),
64 260 : _do_tuning(isParamValid("tune_parameters")),
65 130 : _optimization_opts(StochasticTools::GaussianProcess::GPOptimizerOptions(
66 260 : getParam<unsigned int>("show_every_nth_iteration"),
67 260 : getParam<unsigned int>("num_iters"),
68 260 : getParam<unsigned int>("batch_size"),
69 260 : getParam<Real>("learning_rate"),
70 : 0.9,
71 : 0.999,
72 : 1e-7,
73 : 1e-4,
74 260 : getParam<MooseEnum>("optimizer")
75 : .getEnum<StochasticTools::GaussianProcess::OptimizerType>())),
76 130 : _sampler_row(getSamplerData())
77 : {
78 : // Error Checking
79 260 : if (parameters.isParamSetByUser("batch_size"))
80 66 : if (_sampler.getNumberOfRows() < _optimization_opts.batch_size)
81 2 : paramError("batch_size", "Batch size cannot be greater than the training data set size.");
82 :
83 : std::vector<std::string> tune_parameters(
84 128 : _do_tuning ? getParam<std::vector<std::string>>("tune_parameters")
85 128 : : std::vector<std::string>{});
86 :
87 408 : if (isParamValid("tuning_min") &&
88 176 : (getParam<std::vector<Real>>("tuning_min").size() != tune_parameters.size()))
89 0 : mooseError("tuning_min size does not match tune_parameters");
90 408 : if (isParamValid("tuning_max") &&
91 176 : (getParam<std::vector<Real>>("tuning_max").size() != tune_parameters.size()))
92 0 : mooseError("tuning_max size does not match tune_parameters");
93 :
94 : std::vector<Real> lower_bounds, upper_bounds;
95 256 : if (isParamValid("tuning_min"))
96 72 : lower_bounds = getParam<std::vector<Real>>("tuning_min");
97 256 : if (isParamValid("tuning_max"))
98 72 : upper_bounds = getParam<std::vector<Real>>("tuning_max");
99 :
100 128 : _gp.initialize(getCovarianceFunctionByName(parameters.get<UserObjectName>("covariance_function")),
101 : tune_parameters,
102 : lower_bounds,
103 : upper_bounds);
104 :
105 128 : _n_outputs = _gp.getCovarFunction().numOutputs();
106 128 : }
107 :
108 : void
109 176 : GaussianProcessTrainer::preTrain()
110 : {
111 176 : _params_buffer.clear();
112 176 : _data_buffer.clear();
113 176 : _params_buffer.reserve(getLocalSampleSize());
114 176 : _data_buffer.reserve(getLocalSampleSize());
115 176 : }
116 :
117 : void
118 1560 : GaussianProcessTrainer::train()
119 : {
120 1560 : _params_buffer.push_back(_predictor_row);
121 :
122 1560 : if (_rvecval && _rvecval->size() != _n_outputs)
123 0 : mooseError("The size of the provided response (",
124 0 : _rvecval->size(),
125 : ") does not match the number of expected outputs from the covariance (",
126 : _n_outputs,
127 : ")!");
128 :
129 1560 : _data_buffer.push_back(_rvecval ? (*_rvecval) : std::vector<Real>(1, *_rval));
130 1560 : }
131 :
132 : void
133 176 : GaussianProcessTrainer::postTrain()
134 : {
135 : // Instead of gatherSum, we have to allgather.
136 176 : _communicator.allgather(_params_buffer);
137 176 : _communicator.allgather(_data_buffer);
138 :
139 224 : _training_params = torch::empty({long(_params_buffer.size()), _n_dims}, at::kDouble);
140 176 : _training_data = torch::empty({long(_data_buffer.size()), _n_outputs}, at::kDouble);
141 :
142 176 : auto params_accessor = _training_params.accessor<Real, 2>();
143 176 : auto data_accessor = _training_data.accessor<Real, 2>();
144 :
145 2672 : for (auto ii : make_range(_training_params.sizes()[0]))
146 : {
147 8032 : for (auto jj : make_range(_n_dims))
148 5536 : params_accessor[ii][jj] = _params_buffer[ii][jj];
149 5152 : for (auto jj : make_range(_n_outputs))
150 2656 : data_accessor[ii][jj] = _data_buffer[ii][jj];
151 : }
152 :
153 176 : LibtorchUtils::moveToLibtorchDevice(_training_params, _app.getLibtorchDevice());
154 176 : LibtorchUtils::moveToLibtorchDevice(_training_data, _app.getLibtorchDevice());
155 :
156 : // Standardize (center and scale) training params
157 176 : if (_standardize_params)
158 176 : _gp.standardizeParameters(_training_params);
159 : // if not standardizing data set mean=0, std=1 for use in surrogate
160 : else
161 0 : _gp.paramStandardizer().set(0, 1, _n_dims);
162 : // Standardize (center and scale) training data
163 176 : if (_standardize_data)
164 176 : _gp.standardizeData(_training_data);
165 : // if not standardizing data set mean=0, std=1 for use in surrogate
166 : else
167 0 : _gp.dataStandardizer().set(0, 1, _n_outputs);
168 :
169 : // Setup the covariance
170 176 : _gp.setupCovarianceMatrix(_training_params, _training_data, _optimization_opts);
171 176 : }
172 :
173 : #endif
|