folivo/PINA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refact

Browse Source
This commit is contained in:
Nicola Demo
2024-08-05 17:34:34 +02:00
parent 686b557144
commit 5245a0b68c
19 changed files with 483 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files

255
pina/solvers/solver.py
View File

@@ -5,10 +5,173 @@ from ..model.network import Network
import pytorch_lightning
from ..utils import check_consistency
from ..problem import AbstractProblem
from ..optim import Optimizer, Scheduler
import torch
import sys
# class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
# """
# Solver base class. This class inherits is a wrapper of
# LightningModule class, inheriting all the
# LightningModule methods.
# """
# def __init__(
# self,
# models,
# problem,
# optimizers,
# optimizers_kwargs,
# extra_features=None,
# ):
# """
# :param models: A torch neural network model instance.
# :type models: torch.nn.Module
# :param problem: A problem definition instance.
# :type problem: AbstractProblem
# :param list(torch.optim.Optimizer) optimizer: A list of neural network optimizers to
# use.
# :param list(dict) optimizer_kwargs: A list of optimizer constructor keyword args.
# :param list(torch.nn.Module) extra_features: The additional input
# features to use as augmented input. If ``None`` no extra features
# are passed. If it is a list of :class:`torch.nn.Module`, the extra feature
# list is passed to all models. If it is a list of extra features' lists,
# each single list of extra feature is passed to a model.
# """
# super().__init__()
# # check consistency of the inputs
# check_consistency(models, torch.nn.Module)
# check_consistency(problem, AbstractProblem)
# check_consistency(optimizers, torch.optim.Optimizer, subclass=True)
# check_consistency(optimizers_kwargs, dict)
# # put everything in a list if only one input
# if not isinstance(models, list):
# models = [models]
# if not isinstance(optimizers, list):
# optimizers = [optimizers]
# optimizers_kwargs = [optimizers_kwargs]
# # number of models and optimizers
# len_model = len(models)
# len_optimizer = len(optimizers)
# len_optimizer_kwargs = len(optimizers_kwargs)
# # check length consistency optimizers
# if len_model != len_optimizer:
# raise ValueError(
# "You must define one optimizer for each model."
# f"Got {len_model} models, and {len_optimizer}"
# " optimizers."
# )
# # check length consistency optimizers kwargs
# if len_optimizer_kwargs != len_optimizer:
# raise ValueError(
# "You must define one dictionary of keyword"
# " arguments for each optimizers."
# f"Got {len_optimizer} optimizers, and"
# f" {len_optimizer_kwargs} dicitionaries"
# )
# # extra features handling
# if (extra_features is None) or (len(extra_features) == 0):
# extra_features = [None] * len_model
# else:
# # if we only have a list of extra features
# if not isinstance(extra_features[0], (tuple, list)):
# extra_features = [extra_features] * len_model
# else: # if we have a list of list extra features
# if len(extra_features) != len_model:
# raise ValueError(
# "You passed a list of extrafeatures list with len"
# f"different of models len. Expected {len_model} "
# f"got {len(extra_features)}. If you want to use "
# "the same list of extra features for all models, "
# "just pass a list of extrafeatures and not a list "
# "of list of extra features."
# )
# # assigning model and optimizers
# self._pina_models = []
# self._pina_optimizers = []
# for idx in range(len_model):
# model_ = Network(
# model=models[idx],
# input_variables=problem.input_variables,
# output_variables=problem.output_variables,
# extra_features=extra_features[idx],
# )
# optim_ = optimizers[idx](
# model_.parameters(), **optimizers_kwargs[idx]
# )
# self._pina_models.append(model_)
# self._pina_optimizers.append(optim_)
# # assigning problem
# self._pina_problem = problem
# @abstractmethod
# def forward(self, *args, **kwargs):
# pass
# @abstractmethod
# def training_step(self):
# pass
# @abstractmethod
# def configure_optimizers(self):
# pass
# @property
# def models(self):
# """
# The torch model."""
# return self._pina_models
# @property
# def optimizers(self):
# """
# The torch model."""
# return self._pina_optimizers
# @property
# def problem(self):
# """
# The problem formulation."""
# return self._pina_problem
# def on_train_start(self):
# """
# On training epoch start this function is call to do global checks for
# the different solvers.
# """
# # 1. Check the verison for dataloader
# dataloader = self.trainer.train_dataloader
# if sys.version_info < (3, 8):
# dataloader = dataloader.loaders
# self._dataloader = dataloader
# return super().on_train_start()
# @model.setter
# def model(self, new_model):
# """
# Set the torch."""
# check_consistency(new_model, nn.Module, 'torch model')
# self._model= new_model
# @problem.setter
# def problem(self, problem):
# """
# Set the problem formulation."""
# check_consistency(problem, AbstractProblem, 'pina problem')
# self._problem = problem
class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
"""
Solver base class. This class inherits is a wrapper of
@@ -18,45 +181,36 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
def __init__(
self,
models,
model,
problem,
optimizers,
optimizers_kwargs,
extra_features=None,
optimizer,
scheduler,
):
"""
:param models: A torch neural network model instance.
:type models: torch.nn.Module
:param model: A torch neural network model instance.
:type model: torch.nn.Module
:param problem: A problem definition instance.
:type problem: AbstractProblem
:param list(torch.optim.Optimizer) optimizer: A list of neural network optimizers to
use.
:param list(dict) optimizer_kwargs: A list of optimizer constructor keyword args.
:param list(torch.nn.Module) extra_features: The additional input
features to use as augmented input. If ``None`` no extra features
are passed. If it is a list of :class:`torch.nn.Module`, the extra feature
list is passed to all models. If it is a list of extra features' lists,
each single list of extra feature is passed to a model.
:param list(torch.optim.Optimizer) optimizer: A list of neural network
optimizers to use.
"""
super().__init__()
# check consistency of the inputs
check_consistency(models, torch.nn.Module)
check_consistency(model, torch.nn.Module)
check_consistency(problem, AbstractProblem)
check_consistency(optimizers, torch.optim.Optimizer, subclass=True)
check_consistency(optimizers_kwargs, dict)
check_consistency(optimizer, Optimizer)
check_consistency(scheduler, Scheduler)
# put everything in a list if only one input
if not isinstance(models, list):
models = [models]
if not isinstance(optimizers, list):
optimizers = [optimizers]
optimizers_kwargs = [optimizers_kwargs]
if not isinstance(model, list):
model = [model]
if not isinstance(optimizer, list):
optimizer = [optimizer]
# number of models and optimizers
len_model = len(models)
len_optimizer = len(optimizers)
len_optimizer_kwargs = len(optimizers_kwargs)
len_model = len(model)
len_optimizer = len(optimizer)
# check length consistency optimizers
if len_model != len_optimizer:
@@ -66,52 +220,11 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
" optimizers."
)
# check length consistency optimizers kwargs
if len_optimizer_kwargs != len_optimizer:
raise ValueError(
"You must define one dictionary of keyword"
" arguments for each optimizers."
f"Got {len_optimizer} optimizers, and"
f" {len_optimizer_kwargs} dicitionaries"
)
# extra features handling
if (extra_features is None) or (len(extra_features) == 0):
extra_features = [None] * len_model
else:
# if we only have a list of extra features
if not isinstance(extra_features[0], (tuple, list)):
extra_features = [extra_features] * len_model
else: # if we have a list of list extra features
if len(extra_features) != len_model:
raise ValueError(
"You passed a list of extrafeatures list with len"
f"different of models len. Expected {len_model} "
f"got {len(extra_features)}. If you want to use "
"the same list of extra features for all models, "
"just pass a list of extrafeatures and not a list "
"of list of extra features."
)
# assigning model and optimizers
self._pina_models = []
self._pina_optimizers = []
for idx in range(len_model):
model_ = Network(
model=models[idx],
input_variables=problem.input_variables,
output_variables=problem.output_variables,
extra_features=extra_features[idx],
)
optim_ = optimizers[idx](
model_.parameters(), **optimizers_kwargs[idx]
)
self._pina_models.append(model_)
self._pina_optimizers.append(optim_)
# assigning problem
self._pina_problem = problem
self._pina_model = model
self._pina_optimizer = optimizer
self._pina_scheduler = scheduler
@abstractmethod
def forward(self, *args, **kwargs):
@@ -129,13 +242,13 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
def models(self):
"""
The torch model."""
return self._pina_models
return self._pina_model
@property
def optimizers(self):
"""
The torch model."""
return self._pina_optimizers
return self._pina_optimizer
@property
def problem(self):

63
pina/solvers/supervised.py
View File

@@ -1,21 +1,14 @@
""" Module for SupervisedSolver """
import torch
from torch.nn.modules.loss import _Loss
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 torch.optim.lr_scheduler import ConstantLR
from ..optim import Optimizer, Scheduler, TorchOptimizer, TorchScheduler
from .solver import SolverInterface
from ..label_tensor import LabelTensor
from ..utils import check_consistency
from ..loss import LossInterface
from torch.nn.modules.loss import _Loss
class SupervisedSolver(SolverInterface):
@@ -51,12 +44,9 @@ class SupervisedSolver(SolverInterface):
self,
problem,
model,
extra_features=None,
loss=torch.nn.MSELoss(),
optimizer=torch.optim.Adam,
optimizer_kwargs={"lr": 0.001},
scheduler=ConstantLR,
scheduler_kwargs={"factor": 1, "total_iters": 0},
loss=None,
optimizer=None,
scheduler=None,
):
"""
:param AbstractProblem problem: The formualation of the problem.
@@ -73,24 +63,26 @@ class SupervisedSolver(SolverInterface):
rate scheduler.
:param dict scheduler_kwargs: LR scheduler constructor keyword args.
"""
if loss is None:
loss = torch.nn.MSELoss()
if optimizer is None:
optimizer = TorchOptimizer(torch.optim.Adam, lr=0.001)
if scheduler is None:
scheduler = TorchScheduler(
torch.optim.lr_scheduler.ConstantLR)
super().__init__(
models=[model],
model=model,
problem=problem,
optimizers=[optimizer],
optimizers_kwargs=[optimizer_kwargs],
extra_features=extra_features,
optimizer=optimizer,
scheduler=scheduler,
)
# check consistency
check_consistency(scheduler, LRScheduler, subclass=True)
check_consistency(scheduler_kwargs, dict)
check_consistency(loss, (LossInterface, _Loss), subclass=False)
# assign variables
self._scheduler = scheduler(self.optimizers[0], **scheduler_kwargs)
self._loss = loss
self._neural_net = self.models[0]
def forward(self, x):
"""Forward pass implementation for the solver.
@@ -98,7 +90,7 @@ class SupervisedSolver(SolverInterface):
:return: Solver solution.
:rtype: torch.Tensor
"""
return self.neural_net(x)
return self._pina_model(x)
def configure_optimizers(self):
"""Optimizer configuration for the solver.
@@ -106,7 +98,9 @@ class SupervisedSolver(SolverInterface):
:return: The optimizers and the schedulers
:rtype: tuple(list, list)
"""
return self.optimizers, [self.scheduler]
self._pina_optimizer.hook(self._pina_model.parameters())
self._pina_scheduler.hook(self._pina_optimizer)
return self._pina_optimizer, self._pina_scheduler
def training_step(self, batch, batch_idx):
"""Solver training step.
@@ -168,14 +162,21 @@ class SupervisedSolver(SolverInterface):
"""
Scheduler for training.
"""
return self._scheduler
return self._pina_scheduler
@property
def optimizer(self):
"""
Optimizer for training.
"""
return self._pina_optimizer
@property
def neural_net(self):
def model(self):
"""
Neural network for training.
"""
return self._neural_net
return self._pina_model
@property
def loss(self):