Update solvers (#434)
* Enable DDP training with batch_size=None and add validity check for split sizes * Refactoring SolverInterfaces (#435) * Solver update + weighting * Updating PINN for 0.2 * Modify GAROM + tests * Adding more versatile loggers * Disable compilation when running on Windows * Fix tests --------- Co-authored-by: giovanni <giovanni.canali98@yahoo.it> Co-authored-by: FilippoOlivo <filippo@filippoolivo.com>
This commit is contained in:
Dario Coscia
committed by
Nicola Demo
Nicola Demo
parent
780c4921eb
commit
9cae9a438f
@@ -2,19 +2,12 @@
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
from torch.optim.lr_scheduler import LRScheduler # torch >= 2.0
|
||||
except ImportError:
|
||||
from torch.optim.lr_scheduler import (
|
||||
_LRScheduler as LRScheduler,
|
||||
) # torch < 2.0
|
||||
|
||||
|
||||
from .pinn_interface import PINNInterface
|
||||
from ..solver import SingleSolverInterface
|
||||
from ...problem import InverseProblem
|
||||
|
||||
|
||||
class PINN(PINNInterface):
|
||||
class PINN(PINNInterface, SingleSolverInterface):
|
||||
r"""
|
||||
Physics Informed Neural Network (PINN) solver class.
|
||||
This class implements Physics Informed Neural
|
||||
@@ -41,7 +34,8 @@ class PINN(PINNInterface):
|
||||
\frac{1}{N}\sum_{i=1}^N
|
||||
\mathcal{L}(\mathcal{B}[\mathbf{u}](\mathbf{x}_i))
|
||||
|
||||
where :math:`\mathcal{L}` is a specific loss function, default Mean Square Error:
|
||||
where :math:`\mathcal{L}` is a specific loss function,
|
||||
default Mean Square Error:
|
||||
|
||||
.. math::
|
||||
\mathcal{L}(v) = \| v \|^2_2.
|
||||
@@ -54,54 +48,31 @@ class PINN(PINNInterface):
|
||||
DOI: `10.1038 <https://doi.org/10.1038/s42254-021-00314-5>`_.
|
||||
"""
|
||||
|
||||
__name__ = 'PINN'
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
problem,
|
||||
model,
|
||||
loss=None,
|
||||
optimizer=None,
|
||||
scheduler=None,
|
||||
):
|
||||
def __init__(self,
|
||||
problem,
|
||||
model,
|
||||
optimizer=None,
|
||||
scheduler=None,
|
||||
weighting=None,
|
||||
loss=None):
|
||||
"""
|
||||
:param AbstractProblem problem: The formulation of the problem.
|
||||
:param torch.nn.Module model: The neural network model to use.
|
||||
:param torch.nn.Module loss: The loss function used as minimizer,
|
||||
default :class:`torch.nn.MSELoss`.
|
||||
:param torch.nn.Module extra_features: The additional input
|
||||
features to use as augmented input.
|
||||
:param AbstractProblem problem: The formulation of the problem.
|
||||
:param torch.optim.Optimizer optimizer: The neural network optimizer to
|
||||
use; default is :class:`torch.optim.Adam`.
|
||||
:param dict optimizer_kwargs: Optimizer constructor keyword args.
|
||||
:param torch.optim.LRScheduler scheduler: Learning
|
||||
rate scheduler.
|
||||
:param dict scheduler_kwargs: LR scheduler constructor keyword args.
|
||||
use; default `None`.
|
||||
:param torch.optim.LRScheduler scheduler: Learning rate scheduler;
|
||||
default `None`.
|
||||
:param WeightingInterface weighting: The weighting schema to use;
|
||||
default `None`.
|
||||
:param torch.nn.Module loss: The loss function to be minimized;
|
||||
default `None`.
|
||||
"""
|
||||
super().__init__(
|
||||
models=model,
|
||||
problem=problem,
|
||||
loss=loss,
|
||||
optimizers=optimizer,
|
||||
schedulers=scheduler,
|
||||
)
|
||||
|
||||
# assign variables
|
||||
self._neural_net = self.models[0]
|
||||
|
||||
def forward(self, x):
|
||||
r"""
|
||||
Forward pass implementation for the PINN solver. It returns the function
|
||||
evaluation :math:`\mathbf{u}(\mathbf{x})` at the control points
|
||||
:math:`\mathbf{x}`.
|
||||
|
||||
:param LabelTensor x: Input tensor for the PINN solver. It expects
|
||||
a tensor :math:`N \times D`, where :math:`N` the number of points
|
||||
in the mesh, :math:`D` the dimension of the problem,
|
||||
:return: PINN solution evaluated at contro points.
|
||||
:rtype: LabelTensor
|
||||
"""
|
||||
return self.neural_net(x)
|
||||
super().__init__(model=model,
|
||||
problem=problem,
|
||||
optimizer=optimizer,
|
||||
scheduler=scheduler,
|
||||
weighting=weighting,
|
||||
loss=loss)
|
||||
|
||||
def loss_phys(self, samples, equation):
|
||||
"""
|
||||
@@ -117,46 +88,31 @@ class PINN(PINNInterface):
|
||||
"""
|
||||
residual = self.compute_residual(samples=samples, equation=equation)
|
||||
loss_value = self.loss(
|
||||
torch.zeros_like(residual), residual
|
||||
torch.zeros_like(residual, requires_grad=True), residual
|
||||
)
|
||||
return loss_value
|
||||
|
||||
def configure_optimizers(self):
|
||||
"""
|
||||
Optimizer configuration for the PINN
|
||||
solver.
|
||||
Optimizer configuration for the PINN solver.
|
||||
|
||||
:return: The optimizers and the schedulers
|
||||
:rtype: tuple(list, list)
|
||||
"""
|
||||
# if the problem is an InverseProblem, add the unknown parameters
|
||||
# to the parameters that the optimizer needs to optimize
|
||||
|
||||
|
||||
self._optimizer.hook(self._model.parameters())
|
||||
# If the problem is an InverseProblem, add the unknown parameters
|
||||
# to the parameters to be optimized.
|
||||
self.optimizer.hook(self.model.parameters())
|
||||
if isinstance(self.problem, InverseProblem):
|
||||
self._optimizer.optimizer_instance.add_param_group(
|
||||
{
|
||||
"params": [
|
||||
self._params[var]
|
||||
for var in self.problem.unknown_variables
|
||||
]
|
||||
}
|
||||
)
|
||||
self._scheduler.hook(self._optimizer)
|
||||
return ([self._optimizer.optimizer_instance],
|
||||
[self._scheduler.scheduler_instance])
|
||||
|
||||
@property
|
||||
def scheduler(self):
|
||||
"""
|
||||
Scheduler for the PINN training.
|
||||
"""
|
||||
return self._scheduler
|
||||
|
||||
@property
|
||||
def neural_net(self):
|
||||
"""
|
||||
Neural network for the PINN training.
|
||||
"""
|
||||
return self._neural_net
|
||||
self.optimizer.instance.add_param_group(
|
||||
{
|
||||
"params": [
|
||||
self._params[var]
|
||||
for var in self.problem.unknown_variables
|
||||
]
|
||||
}
|
||||
)
|
||||
self.scheduler.hook(self.optimizer)
|
||||
return (
|
||||
[self.optimizer.instance],
|
||||
[self.scheduler.instance]
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user