(../moose/test/src/libtorch/vectorpostprocessors/LibtorchTorchScriptNeuralNetTest.C)
#ifdef LIBTORCH_ENABLED
#include <torch/torch.h>
#include "LibtorchTorchScriptNeuralNet.h"
#include "LibtorchTorchScriptNeuralNetTest.h"
registerMooseObject("MooseTestApp", LibtorchTorchScriptNeuralNetTest);
InputParameters
LibtorchTorchScriptNeuralNetTest::validParams()
{
InputParameters params = GeneralVectorPostprocessor::validParams();
params.addParam<std::string>(
"filename", "my_net.pt", "The name of the file where the torch script is saved.");
params.addClassDescription("Evaluates a neural network saved from python.");
return params;
}
LibtorchTorchScriptNeuralNetTest::LibtorchTorchScriptNeuralNetTest(const InputParameters & params)
: GeneralVectorPostprocessor(params),
_x_values(declareVector("x_values")),
_y_values(declareVector("y_values")),
_z_values(declareVector("z_values")),
_nn_values(declareVector("nn_values"))
{
std::shared_ptr<Moose::LibtorchNeuralNetBase> nn =
std::make_shared<Moose::LibtorchTorchScriptNeuralNet>(getParam<std::string>("filename"));
torch::Tensor input = at::zeros({1, 3}, at::kDouble);
_x_values.push_back(0.0);
_y_values.push_back(0.0);
_z_values.push_back(0.0);
torch::Tensor prediction = nn->forward(input);
_nn_values.push_back(prediction.item<double>());
for (unsigned int i = 1; i < 3; ++i)
{
torch::Tensor input = torch::add(at::zeros({1, 3}, at::kDouble), i * 0.5);
torch::Tensor prediction = nn->forward(input);
_x_values.push_back(i * 0.5);
_y_values.push_back(i * 0.5);
_z_values.push_back(i * 0.5);
_nn_values.push_back(prediction.item<double>());
}
}
#endif