neml2::LibtorchModel Class Reference

Evaluate a pretrained libtorch model in .pt format, such as a neural network. More...

#include <LibtorchModel.h>

Inheritance diagram for neml2::LibtorchModel:

Public Member Functions
	LibtorchModel (const OptionSet &options)
virtual void	to (const torch::TensorOptions &options) override
	Override the base implementation to additionally send the model loaded from torch script to different device and dtype. More...
virtual void	request_AD () override

Static Public Member Functions
static OptionSet	expected_options ()

Protected Member Functions
virtual void	set_value (bool out, bool dout_din, bool d2out_din2) override

Protected Attributes
std::vector< const Variable< Scalar > * >	_inputs
std::vector< Variable< Scalar > * >	_outputs
Moose::DataFileUtils::Path	_file_path
std::unique_ptr< torch::jit::script::Module >	_surrogate
	We need to use a pointer here because forward is not const qualified. More...

Detailed Description

Evaluate a pretrained libtorch model in .pt format, such as a neural network.

Evaluates models with an arbitrary number of inputs and maps them to an arbitrary number of outputs.

Definition at line 31 of file LibtorchModel.h.

Constructor & Destructor Documentation

LibtorchModel()

neml2::LibtorchModel::LibtorchModel

(

const OptionSet &

options

)

Definition at line 39 of file LibtorchModel.C.

  : Model(options),
    _file_path(Moose::DataFileUtils::getPath(options.get<std::string>("file_path"))),
    _surrogate(std::make_unique<torch::jit::script::Module>(torch::jit::load(_file_path.path)))
{
  // inputs
  for (const auto & fv : options.get<std::vector<VariableName>>("inputs"))
    _inputs.push_back(&declare_input_variable<Scalar>(fv));
  for (const auto & fv : options.get<std::vector<VariableName>>("outputs"))
    _outputs.push_back(&declare_output_variable<Scalar>(fv));
}

Member Function Documentation

expected_options()

OptionSet neml2::LibtorchModel::expected_options ( )

static

Definition at line 26 of file LibtorchModel.C.

{
  auto options = Model::expected_options();
  options.set<std::vector<VariableName>>("inputs");
  options.set<std::vector<VariableName>>("outputs");
  options.set("outputs").doc() = "The scaled neural network output";
  options.set<std::string>("file_path");
  // No jitting :/
  options.set<bool>("jit") = false;
  options.set("jit").suppressed() = true;
  return options;
}

request_AD()

void neml2::LibtorchModel::request_AD ( )

overridevirtual

Definition at line 64 of file LibtorchModel.C.

{
  std::vector<const VariableBase *> inputs;
  for (size_t i = 0; i < _inputs.size(); ++i)
    inputs.push_back(_inputs[i]);

  for (size_t i = 0; i < _outputs.size(); ++i)
    _outputs[i]->request_AD(inputs);
}

set_value()

void neml2::LibtorchModel::set_value ( bool  out, bool  dout_din, bool  d2out_din2  )

overrideprotectedvirtual

Definition at line 75 of file LibtorchModel.C.

{
  if (out)
  {
    std::vector<at::Tensor> values;
    auto first_batch_dim = _inputs[0]->batch_dim();
    for (size_t i = 0; i < _inputs.size(); ++i)
    {
      // assert that all inputs have the same batch dimension
      neml_assert(_inputs[i]->batch_dim() == first_batch_dim);
      values.push_back(_inputs[i]->value());
    }

    auto x = Tensor(torch::transpose(torch::vstack(at::ArrayRef<at::Tensor>(
                                         values.data(), static_cast<int64_t>(values.size()))),
                                     0,
                                     1),
                    _inputs[0]->batch_dim());

    // Feed forward the neural network and process the output
    auto temp = _surrogate->forward({x}).toTensor().squeeze();
    auto y0 =
        (temp.dim() == 1) ? temp.view({temp.size(0), 1}).transpose(0, 1) : temp.transpose(0, 1);

    for (size_t i = 0; i < _outputs.size(); ++i)
      *_outputs[i] = Scalar(y0[i], _inputs[0]->batch_dim());
  }
}

to()

void neml2::LibtorchModel::to ( const torch::TensorOptions &  options )

overridevirtual

Override the base implementation to additionally send the model loaded from torch script to different device and dtype.

Definition at line 52 of file LibtorchModel.C.

{
  Model::to(options);

  if (options.has_device())
    _surrogate->to(options.device());

  if (options.has_dtype())
    _surrogate->to(torch::Dtype(caffe2::typeMetaToScalarType(options.dtype())));
}

Member Data Documentation

_file_path

Moose::DataFileUtils::Path neml2::LibtorchModel::_file_path

protected

Definition at line 53 of file LibtorchModel.h.

_inputs

std::vector<const Variable<Scalar> *> neml2::LibtorchModel::_inputs

protected

Definition at line 50 of file LibtorchModel.h.

Referenced by LibtorchModel(), request_AD(), and set_value().

_outputs

std::vector<Variable<Scalar> *> neml2::LibtorchModel::_outputs

protected

Definition at line 52 of file LibtorchModel.h.

Referenced by LibtorchModel(), request_AD(), and set_value().

_surrogate

std::unique_ptr<torch::jit::script::Module> neml2::LibtorchModel::_surrogate

protected

We need to use a pointer here because forward is not const qualified.

Definition at line 55 of file LibtorchModel.h.

Referenced by set_value(), and to().

---

The documentation for this class was generated from the following files:

• include/neml2/models/LibtorchModel.h
• src/neml2/models/LibtorchModel.C