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
@@ -1,167 +1,177 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from pina.problem import AbstractProblem
|
||||
import pytest
|
||||
from pina import Condition, LabelTensor
|
||||
from pina.solvers import GAROM
|
||||
from pina.condition import InputOutputPointsCondition
|
||||
from pina.problem import AbstractProblem
|
||||
from pina.model import FeedForward
|
||||
from pina.trainer import Trainer
|
||||
import torch.nn as nn
|
||||
import matplotlib.tri as tri
|
||||
from torch._dynamo.eval_frame import OptimizedModule
|
||||
|
||||
|
||||
# def func(x, mu1, mu2):
|
||||
# import torch
|
||||
# x_m1 = (x[:, 0] - mu1).pow(2)
|
||||
# x_m2 = (x[:, 1] - mu2).pow(2)
|
||||
# norm = x[:, 0]**2 + x[:, 1]**2
|
||||
# return torch.exp(-(x_m1 + x_m2))
|
||||
class TensorProblem(AbstractProblem):
|
||||
input_variables = ['u_0', 'u_1']
|
||||
output_variables = ['u']
|
||||
conditions = {
|
||||
'data': Condition(
|
||||
output_points=torch.randn(50, 2),
|
||||
input_points=torch.randn(50, 1))
|
||||
}
|
||||
|
||||
|
||||
# class ParametricGaussian(AbstractProblem):
|
||||
# output_variables = [f'u_{i}' for i in range(900)]
|
||||
# simple Generator Network
|
||||
class Generator(nn.Module):
|
||||
|
||||
# # params
|
||||
# xx = torch.linspace(-1, 1, 20)
|
||||
# yy = xx
|
||||
# params = LabelTensor(torch.cartesian_prod(xx, yy), labels=['mu1', 'mu2'])
|
||||
def __init__(self,
|
||||
input_dimension=2,
|
||||
parameters_dimension=1,
|
||||
noise_dimension=2,
|
||||
activation=torch.nn.SiLU):
|
||||
super().__init__()
|
||||
|
||||
# # define domain
|
||||
# x = torch.linspace(-1, 1, 30)
|
||||
# domain = torch.cartesian_prod(x, x)
|
||||
# triang = tri.Triangulation(domain[:, 0], domain[:, 1])
|
||||
# sol = []
|
||||
# for p in params:
|
||||
# sol.append(func(domain, p[0], p[1]))
|
||||
# snapshots = LabelTensor(torch.stack(sol), labels=output_variables)
|
||||
self._noise_dimension = noise_dimension
|
||||
self._activation = activation
|
||||
self.model = FeedForward(6*noise_dimension, input_dimension)
|
||||
self.condition = FeedForward(parameters_dimension, 5 * noise_dimension)
|
||||
|
||||
# # define conditions
|
||||
# conditions = {
|
||||
# 'data': Condition(input_points=params, output_points=snapshots)
|
||||
# }
|
||||
def forward(self, param):
|
||||
# uniform sampling in [-1, 1]
|
||||
z = 2 * torch.rand(size=(param.shape[0], self._noise_dimension),
|
||||
device=param.device,
|
||||
dtype=param.dtype,
|
||||
requires_grad=True) - 1
|
||||
return self.model(torch.cat((z, self.condition(param)), dim=-1))
|
||||
|
||||
# Simple Discriminator Network
|
||||
|
||||
|
||||
# # simple Generator Network
|
||||
# class Generator(nn.Module):
|
||||
class Discriminator(nn.Module):
|
||||
|
||||
# def __init__(self,
|
||||
# input_dimension,
|
||||
# parameters_dimension,
|
||||
# noise_dimension,
|
||||
# activation=torch.nn.SiLU):
|
||||
# super().__init__()
|
||||
def __init__(self,
|
||||
input_dimension=2,
|
||||
parameter_dimension=1,
|
||||
hidden_dimension=2,
|
||||
activation=torch.nn.ReLU):
|
||||
super().__init__()
|
||||
|
||||
# self._noise_dimension = noise_dimension
|
||||
# self._activation = activation
|
||||
self._activation = activation
|
||||
self.encoding = FeedForward(input_dimension, hidden_dimension)
|
||||
self.decoding = FeedForward(2*hidden_dimension, input_dimension)
|
||||
self.condition = FeedForward(parameter_dimension, hidden_dimension)
|
||||
|
||||
# self.model = torch.nn.Sequential(
|
||||
# torch.nn.Linear(6 * self._noise_dimension, input_dimension // 6),
|
||||
# self._activation(),
|
||||
# torch.nn.Linear(input_dimension // 6, input_dimension // 3),
|
||||
# self._activation(),
|
||||
# torch.nn.Linear(input_dimension // 3, input_dimension))
|
||||
# self.condition = torch.nn.Sequential(
|
||||
# torch.nn.Linear(parameters_dimension, 2 * self._noise_dimension),
|
||||
# self._activation(),
|
||||
# torch.nn.Linear(2 * self._noise_dimension,
|
||||
# 5 * self._noise_dimension))
|
||||
|
||||
# def forward(self, param):
|
||||
# # uniform sampling in [-1, 1]
|
||||
# z = torch.rand(size=(param.shape[0], self._noise_dimension),
|
||||
# device=param.device,
|
||||
# dtype=param.dtype,
|
||||
# requires_grad=True)
|
||||
# z = 2. * z - 1.
|
||||
|
||||
# # conditioning by concatenation of mapped parameters
|
||||
# input_ = torch.cat((z, self.condition(param)), dim=-1)
|
||||
# out = self.model(input_)
|
||||
|
||||
# return out
|
||||
def forward(self, data):
|
||||
x, condition = data
|
||||
encoding = self.encoding(x)
|
||||
conditioning = torch.cat((encoding, self.condition(condition)), dim=-1)
|
||||
decoding = self.decoding(conditioning)
|
||||
return decoding
|
||||
|
||||
|
||||
# # Simple Discriminator Network
|
||||
# class Discriminator(nn.Module):
|
||||
|
||||
# def __init__(self,
|
||||
# input_dimension,
|
||||
# parameter_dimension,
|
||||
# hidden_dimension,
|
||||
# activation=torch.nn.ReLU):
|
||||
# super().__init__()
|
||||
|
||||
# self._activation = activation
|
||||
# self.encoding = torch.nn.Sequential(
|
||||
# torch.nn.Linear(input_dimension, input_dimension // 3),
|
||||
# self._activation(),
|
||||
# torch.nn.Linear(input_dimension // 3, input_dimension // 6),
|
||||
# self._activation(),
|
||||
# torch.nn.Linear(input_dimension // 6, hidden_dimension))
|
||||
# self.decoding = torch.nn.Sequential(
|
||||
# torch.nn.Linear(2 * hidden_dimension, input_dimension // 6),
|
||||
# self._activation(),
|
||||
# torch.nn.Linear(input_dimension // 6, input_dimension // 3),
|
||||
# self._activation(),
|
||||
# torch.nn.Linear(input_dimension // 3, input_dimension),
|
||||
# )
|
||||
|
||||
# self.condition = torch.nn.Sequential(
|
||||
# torch.nn.Linear(parameter_dimension, hidden_dimension // 2),
|
||||
# self._activation(),
|
||||
# torch.nn.Linear(hidden_dimension // 2, hidden_dimension))
|
||||
|
||||
# def forward(self, data):
|
||||
# x, condition = data
|
||||
# encoding = self.encoding(x)
|
||||
# conditioning = torch.cat((encoding, self.condition(condition)), dim=-1)
|
||||
# decoding = self.decoding(conditioning)
|
||||
# return decoding
|
||||
def test_constructor():
|
||||
GAROM(problem=TensorProblem(),
|
||||
generator=Generator(),
|
||||
discriminator=Discriminator())
|
||||
assert GAROM.accepted_conditions_types == (
|
||||
InputOutputPointsCondition
|
||||
)
|
||||
|
||||
|
||||
# problem = ParametricGaussian()
|
||||
@pytest.mark.parametrize("batch_size", [None, 1, 5, 20])
|
||||
@pytest.mark.parametrize("compile", [True, False])
|
||||
def test_solver_train(batch_size, compile):
|
||||
solver = GAROM(problem=TensorProblem(),
|
||||
generator=Generator(),
|
||||
discriminator=Discriminator())
|
||||
trainer = Trainer(solver=solver,
|
||||
max_epochs=2,
|
||||
accelerator='cpu',
|
||||
batch_size=batch_size,
|
||||
train_size=1.,
|
||||
test_size=0.,
|
||||
val_size=0.,
|
||||
compile=compile)
|
||||
trainer.train()
|
||||
if trainer.compile:
|
||||
assert (all([isinstance(model, OptimizedModule)
|
||||
for model in solver.models]))
|
||||
|
||||
|
||||
# def test_constructor():
|
||||
# GAROM(problem=problem,
|
||||
# generator=Generator(input_dimension=900,
|
||||
# parameters_dimension=2,
|
||||
# noise_dimension=12),
|
||||
# discriminator=Discriminator(input_dimension=900,
|
||||
# parameter_dimension=2,
|
||||
# hidden_dimension=64))
|
||||
@pytest.mark.parametrize("batch_size", [None, 1, 5, 20])
|
||||
@pytest.mark.parametrize("compile", [True, False])
|
||||
def test_solver_validation(batch_size, compile):
|
||||
solver = GAROM(problem=TensorProblem(),
|
||||
generator=Generator(),
|
||||
discriminator=Discriminator())
|
||||
|
||||
trainer = Trainer(solver=solver,
|
||||
max_epochs=2,
|
||||
accelerator='cpu',
|
||||
batch_size=batch_size,
|
||||
train_size=0.9,
|
||||
val_size=0.1,
|
||||
test_size=0.,
|
||||
compile=compile)
|
||||
trainer.train()
|
||||
if trainer.compile:
|
||||
assert (all([isinstance(model, OptimizedModule)
|
||||
for model in solver.models]))
|
||||
|
||||
|
||||
# def test_train_cpu():
|
||||
# solver = GAROM(problem=problem,
|
||||
# generator=Generator(input_dimension=900,
|
||||
# parameters_dimension=2,
|
||||
# noise_dimension=12),
|
||||
# discriminator=Discriminator(input_dimension=900,
|
||||
# parameter_dimension=2,
|
||||
# hidden_dimension=64))
|
||||
|
||||
# trainer = Trainer(solver=solver, max_epochs=4, accelerator='cpu', batch_size=20)
|
||||
# trainer.train()
|
||||
@pytest.mark.parametrize("batch_size", [None, 1, 5, 20])
|
||||
@pytest.mark.parametrize("compile", [True, False])
|
||||
def test_solver_test(batch_size, compile):
|
||||
solver = GAROM(problem=TensorProblem(),
|
||||
generator=Generator(),
|
||||
discriminator=Discriminator(),
|
||||
)
|
||||
trainer = Trainer(solver=solver,
|
||||
max_epochs=2,
|
||||
accelerator='cpu',
|
||||
batch_size=batch_size,
|
||||
train_size=0.8,
|
||||
val_size=0.1,
|
||||
test_size=0.1,
|
||||
compile=compile)
|
||||
trainer.test()
|
||||
if trainer.compile:
|
||||
assert (all([isinstance(model, OptimizedModule)
|
||||
for model in solver.models]))
|
||||
|
||||
|
||||
# def test_sample():
|
||||
# solver = GAROM(problem=problem,
|
||||
# generator=Generator(input_dimension=900,
|
||||
# parameters_dimension=2,
|
||||
# noise_dimension=12),
|
||||
# discriminator=Discriminator(input_dimension=900,
|
||||
# parameter_dimension=2,
|
||||
# hidden_dimension=64))
|
||||
# solver.sample(problem.params)
|
||||
# assert solver.sample(problem.params).shape == problem.snapshots.shape
|
||||
def test_train_load_restore():
|
||||
dir = "tests/test_solvers/tmp/"
|
||||
problem = TensorProblem()
|
||||
solver = GAROM(problem=TensorProblem(),
|
||||
generator=Generator(),
|
||||
discriminator=Discriminator(),
|
||||
)
|
||||
trainer = Trainer(solver=solver,
|
||||
max_epochs=5,
|
||||
accelerator='cpu',
|
||||
batch_size=None,
|
||||
train_size=0.9,
|
||||
test_size=0.1,
|
||||
val_size=0.,
|
||||
default_root_dir=dir)
|
||||
trainer.train()
|
||||
|
||||
# restore
|
||||
new_trainer = Trainer(solver=solver, max_epochs=5, accelerator='cpu')
|
||||
new_trainer.train(
|
||||
ckpt_path=f'{dir}/lightning_logs/version_0/checkpoints/' +
|
||||
'epoch=4-step=5.ckpt')
|
||||
|
||||
# def test_forward():
|
||||
# solver = GAROM(problem=problem,
|
||||
# generator=Generator(input_dimension=900,
|
||||
# parameters_dimension=2,
|
||||
# noise_dimension=12),
|
||||
# discriminator=Discriminator(input_dimension=900,
|
||||
# parameter_dimension=2,
|
||||
# hidden_dimension=64))
|
||||
# solver(problem.params, mc_steps=100, variance=True)
|
||||
# assert solver(problem.params).shape == problem.snapshots.shape
|
||||
# loading
|
||||
new_solver = GAROM.load_from_checkpoint(
|
||||
f'{dir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt',
|
||||
problem=TensorProblem(), generator=Generator(), discriminator=Discriminator())
|
||||
|
||||
test_pts = torch.rand(20, 1)
|
||||
assert new_solver.forward(test_pts).shape == (20, 2)
|
||||
assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
|
||||
|
||||
# rm directories
|
||||
import shutil
|
||||
shutil.rmtree('tests/test_solvers/tmp')
|
||||
|
||||
Reference in New Issue
Block a user