firedrake.ml.pytorch package¶
Submodules¶
firedrake.ml.pytorch.fem_operator module¶
- class firedrake.ml.pytorch.fem_operator.FiredrakeTorchOperator[source]¶
Bases:
Function
PyTorch custom operator representing a set of Firedrake operations expressed as a reduced functional F.
FiredrakeTorchOperator is a wrapper around
torch.autograd.Function
that executes forward and backward passes by directly calling the reduced functional F.- Parameters:
metadata (dict) – Dictionary used to stash Firedrake objects.
*x_P (torch.Tensor) – PyTorch tensors representing the inputs to the Firedrake operator F.
- Returns:
PyTorch tensor representing the output of the Firedrake operator F.
- Return type:
- firedrake.ml.pytorch.fem_operator.fem_operator(F)[source]¶
Cast a Firedrake reduced functional to a PyTorch operator.
The resulting
FiredrakeTorchOperator
will take PyTorch tensors as inputs and return PyTorch tensors as outputs.- Parameters:
F (pyadjoint.ReducedFunctional) – The reduced functional to wrap.
- Returns:
A PyTorch custom operator that wraps the reduced functional F.
- Return type:
- firedrake.ml.pytorch.fem_operator.from_torch(x, V=None)[source]¶
Convert a PyTorch tensor x into a Firedrake object.
- Parameters:
x (torch.Tensor) – PyTorch tensor to convert.
V (firedrake.functionspaceimpl.WithGeometry or None) – Function space of the corresponding
Function
or None when x is to be mapped to aConstant
.
- Returns:
Firedrake object representing the PyTorch tensor x.
- Return type:
- firedrake.ml.pytorch.fem_operator.to_torch(x, gather=False, batched=True, **kwargs)[source]¶
Convert a Firedrake object x into a PyTorch tensor.
- Parameters:
x (firedrake.function.Function, firedrake.vector.Vector or firedrake.constant.Constant) – Firedrake object to convert.
gather (bool) – If True, gather data from all processes
batched (bool) – If True, add a batch dimension to the tensor
kwargs (dict) –
- Additional arguments to be passed to the
torch.Tensor
constructor such as: device: device on which the tensor is allocated (default: “cpu”)
dtype: the desired data type of returned tensor (default: type of x.dat.data)
requires_grad: if the tensor should be annotated (default: False)
- Additional arguments to be passed to the
- Returns:
PyTorch tensor representing the Firedrake object x.
- Return type:
- firedrake.ml.pytorch.fem_operator.torch_operator(F)[source]¶
Cast a Firedrake reduced functional to a PyTorch operator.
The resulting
FiredrakeTorchOperator
will take PyTorch tensors as inputs and return PyTorch tensors as outputs.- Parameters:
F (pyadjoint.ReducedFunctional) – The reduced functional to wrap.
- Returns:
A PyTorch custom operator that wraps the reduced functional F.
- Return type:
firedrake.ml.pytorch.ml_operator module¶
- class firedrake.ml.pytorch.ml_operator.PytorchOperator(*operands, function_space, derivatives=None, argument_slots=(), operator_data)[source]¶
Bases:
MLOperator
External operator class representing machine learning models implemented in PyTorch.
The
PytorchOperator
allows users to embed machine learning models implemented in PyTorch into PDE systems implemented in Firedrake. The actual evaluation of thePytorchOperator
is delegated to the specified PyTorch model. Similarly, differentiation through thePytorchOperator
class is achieved via the torch.autograd module, which provides automatic differentiation capabilities that can be applied on the PyTorch model associated with thePytorchOperator
object.- Parameters:
*operands (ufl.core.expr.Expr or ufl.form.BaseForm) – Operands of the
PytorchOperator
.function_space (firedrake.functionspaceimpl.WithGeometryBase) – The function space the ML operator is mapping to.
derivatives (tuple) – Tuple specifiying the derivative multiindex.
*argument_slots (ufl.coefficient.BaseCoefficient or ufl.argument.BaseArgument) – Tuple containing the arguments of the linear form associated with the ML operator, i.e. the arguments with respect to which the ML operator is linear. Those arguments can be ufl.Argument objects, as a result of differentiation, or ufl.Coefficient objects, as a result of taking the action on a given function.
operator_data (dict) – Dictionary to stash external data specific to the ML operator. This dictionary must at least contain the following: (i) ‘model’: The machine learning model implemented in PyTorch. (ii) ‘inputs_format’: The format of the inputs to the ML model: 0 for models acting globally on the inputs, 1 when acting locally/pointwise on the inputs. Other strategies can also be considered by subclassing the
PytorchOperator
class.
- property model_output¶
- ufl_operands¶
- firedrake.ml.pytorch.ml_operator.ml_operator(model, function_space, inputs_format=0)[source]¶
Helper function for instantiating the
PytorchOperator
class.This function facilitates having a two-stage instantiation which dissociates between class arguments that are fixed, such as the function space or the ML model, and the operands of the operator, which may change, e.g. when the operator is used in a time-loop.
Example
# Stage 1: Partially initialise the operator. N = ml_operator(model, function_space=V) # Stage 2: Define the operands and use the operator in a UFL expression. F = (inner(grad(u), grad(v)) + inner(N(u), v) - inner(f, v)) * dx
- Parameters:
model (collections.abc.Callable) – The PyTorch model to embed in Firedrake.
function_space (firedrake.functionspaceimpl.WithGeometryBase) – The function space into which the machine learning model is mapping.
inputs_format (int) – The format of the input data of the ML model: 0 for models acting globally on the inputs, 1 when acting locally/pointwise on the inputs. Other strategies can also be considered by subclassing the
PytorchOperator
class.
- Returns:
The partially initialised
PytorchOperator
class.- Return type:
- firedrake.ml.pytorch.ml_operator.neuralnet(model, function_space, inputs_format=0)[source]¶
Helper function for instantiating the
PytorchOperator
class.This function facilitates having a two-stage instantiation which dissociates between class arguments that are fixed, such as the function space or the ML model, and the operands of the operator, which may change, e.g. when the operator is used in a time-loop.
Example
# Stage 1: Partially initialise the operator. N = ml_operator(model, function_space=V) # Stage 2: Define the operands and use the operator in a UFL expression. F = (inner(grad(u), grad(v)) + inner(N(u), v) - inner(f, v)) * dx
- Parameters:
model (collections.abc.Callable) – The PyTorch model to embed in Firedrake.
function_space (firedrake.functionspaceimpl.WithGeometryBase) – The function space into which the machine learning model is mapping.
inputs_format (int) – The format of the input data of the ML model: 0 for models acting globally on the inputs, 1 when acting locally/pointwise on the inputs. Other strategies can also be considered by subclassing the
PytorchOperator
class.
- Returns:
The partially initialised
PytorchOperator
class.- Return type: