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 "LibtorchNeuralNetControl.h"
13 : #include "TorchScriptModule.h"
14 : #include "LibtorchUtils.h"
15 :
16 : #include "Transient.h"
17 :
18 : registerMooseObject("MooseApp", LibtorchNeuralNetControl);
19 :
20 : InputParameters
21 1568 : LibtorchNeuralNetControl::validParams()
22 : {
23 1568 : InputParameters params = Control::validParams();
24 1568 : params.addClassDescription("Controls the value of multiple controllable input parameters using a "
25 : "Libtorch-based neural network.");
26 1568 : params.addRequiredParam<std::vector<std::string>>("parameters",
27 : "The input parameter(s) to control.");
28 1568 : params.addRequiredParam<std::vector<PostprocessorName>>(
29 : "responses", "The responses (prostprocessors) which are used for the control.");
30 1568 : params.addParam<std::vector<Real>>(
31 : "response_shift_factors",
32 : "Constants which will be used to shift the response values. This is used for the "
33 : "manipulation of the neural net inputs for better training efficiency.");
34 1568 : params.addParam<std::vector<Real>>(
35 : "response_scaling_factors",
36 : "Constants which will be used to multiply the shifted response values. This is used for "
37 : "the manipulation of the neural net inputs for better training efficiency.");
38 1568 : params.addParam<std::string>("filename",
39 : "Define if the neural net is supposed to be loaded from a file.");
40 4704 : params.addParam<bool>("torch_script_format",
41 3136 : false,
42 : "If we want to load the neural net using the torch-script format.");
43 4704 : params.addParam<unsigned int>(
44 : "input_timesteps",
45 3136 : 1,
46 : "Number of time steps to use in the input data, if larger than 1, "
47 : "data from the previous timesteps will be used as well as inputs in the training.");
48 1568 : params.addParam<std::vector<unsigned int>>("num_neurons_per_layer",
49 : "The number of neurons on each hidden layer.");
50 4704 : params.addParam<std::vector<std::string>>(
51 : "activation_function",
52 6272 : std::vector<std::string>({"relu"}),
53 : "The type of activation functions to use. It is either one value "
54 : "or one value per hidden layer.");
55 :
56 1568 : params.addParam<std::vector<Real>>(
57 : "action_scaling_factors",
58 : "Scale factor that multiplies the NN output to obtain a physically meaningful value.");
59 :
60 1568 : return params;
61 3136 : }
62 :
63 4 : LibtorchNeuralNetControl::LibtorchNeuralNetControl(const InputParameters & parameters)
64 : : Control(parameters),
65 4 : _old_responses(declareRestartableData<std::vector<std::vector<Real>>>("old_responses")),
66 4 : _control_names(getParam<std::vector<std::string>>("parameters")),
67 4 : _current_control_signals(std::vector<Real>(_control_names.size(), 0.0)),
68 4 : _response_names(getParam<std::vector<PostprocessorName>>("responses")),
69 4 : _input_timesteps(getParam<unsigned int>("input_timesteps")),
70 16 : _response_shift_factors(isParamValid("response_shift_factors")
71 4 : ? getParam<std::vector<Real>>("response_shift_factors")
72 8 : : std::vector<Real>(_response_names.size(), 0.0)),
73 16 : _response_scaling_factors(isParamValid("response_scaling_factors")
74 4 : ? getParam<std::vector<Real>>("response_scaling_factors")
75 8 : : std::vector<Real>(_response_names.size(), 1.0)),
76 16 : _action_scaling_factors(isParamValid("action_scaling_factors")
77 4 : ? getParam<std::vector<Real>>("action_scaling_factors")
78 20 : : std::vector<Real>(_control_names.size(), 1.0))
79 : {
80 : // We first check if the input parameters make sense and throw errors if different parameter
81 : // combinations are not allowed
82 10 : conditionalParameterError("filename",
83 : {"num_neurons_per_layer", "activation_function"},
84 7 : !getParam<bool>("torch_script_format"));
85 :
86 3 : if (_response_names.size() != _response_shift_factors.size())
87 0 : paramError("response_shift_factors",
88 : "The number of shift factors is not the same as the number of responses!");
89 :
90 3 : if (_response_names.size() != _response_scaling_factors.size())
91 0 : paramError(
92 : "response_scaling_factors",
93 : "The number of normalization coefficients is not the same as the number of responses!");
94 :
95 3 : if (_control_names.size() != _action_scaling_factors.size())
96 0 : paramError("action_scaling_factors",
97 : "The number of normalization coefficients is not the same as the number of "
98 : "controlled parameters!");
99 :
100 : // We link to the postprocessor values so that we can fetch them any time. This also raises
101 : // errors if we don't have the postprocessors requested in the input.
102 6 : for (unsigned int resp_i = 0; resp_i < _response_names.size(); ++resp_i)
103 3 : _response_values.push_back(&getPostprocessorValueByName(_response_names[resp_i]));
104 :
105 : // If the user wants to read the neural net from file, we do it. We can read it from a
106 : // torchscript file, or we can create a shell and read back the parameters
107 3 : if (isParamValid("filename"))
108 : {
109 2 : std::string filename = getParam<std::string>("filename");
110 2 : if (getParam<bool>("torch_script_format"))
111 1 : _nn = std::make_shared<Moose::TorchScriptModule>(filename);
112 : else
113 : {
114 1 : unsigned int num_inputs = _response_names.size() * _input_timesteps;
115 1 : unsigned int num_outputs = _control_names.size();
116 : std::vector<unsigned int> num_neurons_per_layer =
117 1 : getParam<std::vector<unsigned int>>("num_neurons_per_layer");
118 : std::vector<std::string> activation_functions =
119 2 : parameters.isParamSetByUser("activation_function")
120 1 : ? getParam<std::vector<std::string>>("activation_function")
121 2 : : std::vector<std::string>({"relu"});
122 : auto nn = std::make_shared<Moose::LibtorchArtificialNeuralNet>(
123 1 : filename, num_inputs, num_outputs, num_neurons_per_layer, activation_functions);
124 :
125 : try
126 : {
127 1 : torch::load(nn, filename);
128 1 : _nn = std::make_shared<Moose::LibtorchArtificialNeuralNet>(*nn);
129 : }
130 0 : catch (const c10::Error & e)
131 : {
132 0 : mooseError(
133 : "The requested pytorch parameter file could not be loaded. This can either be the"
134 : "result of the file not existing or a misalignment in the generated container and"
135 : "the data in the file. Make sure the dimensions of the generated neural net are the"
136 : "same as the dimensions of the parameters in the input file!\n",
137 0 : e.msg());
138 0 : }
139 1 : }
140 2 : }
141 11 : }
142 :
143 : void
144 75 : LibtorchNeuralNetControl::execute()
145 : {
146 75 : if (_nn)
147 : {
148 50 : const unsigned int n_controls = _control_names.size();
149 50 : const unsigned int num_old_timesteps = _input_timesteps - 1;
150 :
151 : // Fetch current reporter values and populate _current_response
152 50 : updateCurrentResponse();
153 :
154 : // If this is the first timestep, we fill up the old values with the initial value
155 50 : if (_old_responses.empty())
156 2 : _old_responses.assign(num_old_timesteps, _current_response);
157 :
158 : // Organize the old an current solution into a tensor so we can evaluate the neural net
159 50 : torch::Tensor input_tensor = prepareInputTensor();
160 :
161 : // Evaluate the neural network to get the control values then convert it back to vectors
162 50 : torch::Tensor action = _nn->forward(input_tensor);
163 :
164 50 : _current_control_signals = {action.data_ptr<Real>(), action.data_ptr<Real>() + action.size(1)};
165 100 : for (unsigned int control_i = 0; control_i < n_controls; ++control_i)
166 : {
167 : // We scale the controllable value for physically meaningful control action
168 50 : setControllableValueByName<Real>(_control_names[control_i],
169 150 : _current_control_signals[control_i] *
170 50 : _action_scaling_factors[control_i]);
171 : }
172 :
173 : // We add the curent solution to the old solutions and move everything in there one step
174 : // backward
175 50 : std::rotate(_old_responses.rbegin(), _old_responses.rbegin() + 1, _old_responses.rend());
176 50 : _old_responses[0] = _current_response;
177 50 : }
178 75 : }
179 :
180 : Real
181 78 : LibtorchNeuralNetControl::getSignal(const unsigned int signal_index) const
182 : {
183 : mooseAssert(signal_index < _control_names.size(),
184 : "The index of the requested control signal is not in the [0," +
185 : std::to_string(_control_names.size()) + ") range!");
186 78 : return _current_control_signals[signal_index];
187 : }
188 :
189 : void
190 4 : LibtorchNeuralNetControl::conditionalParameterError(
191 : const std::string & param_name,
192 : const std::vector<std::string> & conditional_params,
193 : bool should_be_defined)
194 : {
195 4 : if (parameters().isParamSetByUser(param_name))
196 7 : for (const auto & param : conditional_params)
197 5 : if (parameters().isParamSetByUser(param) != should_be_defined)
198 1 : paramError(param,
199 : "This parameter should",
200 : (should_be_defined ? " " : " not "),
201 : "be defined when ",
202 : param_name,
203 : " is defined!");
204 3 : }
205 :
206 : void
207 50 : LibtorchNeuralNetControl::updateCurrentResponse()
208 : {
209 : // Gather the current response values from the reporters
210 50 : _current_response.clear();
211 :
212 100 : for (const auto & resp_i : index_range(_response_names))
213 100 : _current_response.push_back((*_response_values[resp_i] - _response_shift_factors[resp_i]) *
214 50 : _response_scaling_factors[resp_i]);
215 50 : }
216 :
217 : void
218 0 : LibtorchNeuralNetControl::loadControlNeuralNet(const Moose::LibtorchArtificialNeuralNet & input_nn)
219 : {
220 0 : _nn = std::make_shared<Moose::LibtorchArtificialNeuralNet>(input_nn);
221 0 : }
222 :
223 : torch::Tensor
224 50 : LibtorchNeuralNetControl::prepareInputTensor()
225 : {
226 50 : const unsigned int num_old_timesteps = _input_timesteps - 1;
227 :
228 : // We convert the standard vectors to libtorch tensors
229 50 : std::vector<Real> raw_input(_current_response);
230 :
231 100 : for (const auto & step_i : make_range(num_old_timesteps))
232 50 : raw_input.insert(raw_input.end(), _old_responses[step_i].begin(), _old_responses[step_i].end());
233 :
234 50 : torch::Tensor input_tensor;
235 50 : LibtorchUtils::vectorToTensor(raw_input, input_tensor);
236 :
237 100 : return input_tensor.transpose(0, 1);
238 50 : }
239 :
240 : const Moose::LibtorchNeuralNetBase &
241 1 : LibtorchNeuralNetControl::controlNeuralNet() const
242 : {
243 1 : if (!hasControlNeuralNet())
244 0 : mooseError("The neural network in the controller must exist!");
245 1 : return *_nn;
246 : }
247 :
248 : #endif
|
