Add Normalizer Callback (#631)
* add normalizer callback * implement shift and scale parameters computation * change name files normalizer data callback * reduce tests * fix documentation * add NotImplementedError for PinaGraphDataset --------- Co-authored-by: FilippoOlivo <filippo@filippoolivo.com> Co-authored-by: giovanni <giovanni.canali98@yahoo.it>
This commit is contained in:
Dario Coscia
244
tests/test_callback/test_normalizer_data_callback.py
Normal file
244
tests/test_callback/test_normalizer_data_callback.py
Normal file
@@ -0,0 +1,244 @@
|
||||
import torch
|
||||
import pytest
|
||||
from copy import deepcopy
|
||||
|
||||
from pina import Trainer, LabelTensor, Condition
|
||||
from pina.solver import SupervisedSolver
|
||||
from pina.model import FeedForward
|
||||
from pina.callback import NormalizerDataCallback
|
||||
from pina.problem import AbstractProblem
|
||||
from pina.problem.zoo import Poisson2DSquareProblem as Poisson
|
||||
from pina.solver import PINN
|
||||
from pina.graph import RadiusGraph
|
||||
|
||||
# for checking normalization
|
||||
stage_map = {
|
||||
"train": ["train_dataset"],
|
||||
"validate": ["val_dataset"],
|
||||
"test": ["test_dataset"],
|
||||
"all": ["train_dataset", "val_dataset", "test_dataset"],
|
||||
}
|
||||
|
||||
input_1 = torch.rand(20, 2) * 10
|
||||
target_1 = torch.rand(20, 1) * 10
|
||||
input_2 = torch.rand(20, 2) * 5
|
||||
target_2 = torch.rand(20, 1) * 5
|
||||
|
||||
|
||||
class LabelTensorProblem(AbstractProblem):
|
||||
input_variables = ["u_0", "u_1"]
|
||||
output_variables = ["u"]
|
||||
conditions = {
|
||||
"data1": Condition(
|
||||
input=LabelTensor(input_1, ["u_0", "u_1"]),
|
||||
target=LabelTensor(target_1, ["u"]),
|
||||
),
|
||||
"data2": Condition(
|
||||
input=LabelTensor(input_2, ["u_0", "u_1"]),
|
||||
target=LabelTensor(target_2, ["u"]),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
class TensorProblem(AbstractProblem):
|
||||
input_variables = ["u_0", "u_1"]
|
||||
output_variables = ["u"]
|
||||
conditions = {
|
||||
"data1": Condition(input=input_1, target=target_1),
|
||||
"data2": Condition(input=input_2, target=target_2),
|
||||
}
|
||||
|
||||
|
||||
input_graph = [RadiusGraph(radius=0.5, pos=torch.rand(10, 2)) for _ in range(5)]
|
||||
output_graph = torch.rand(5, 1)
|
||||
|
||||
|
||||
class GraphProblem(AbstractProblem):
|
||||
input_variables = ["u_0", "u_1"]
|
||||
output_variables = ["u"]
|
||||
conditions = {
|
||||
"data": Condition(input=input_graph, target=output_graph),
|
||||
}
|
||||
|
||||
|
||||
supervised_solver_no_lt = SupervisedSolver(
|
||||
problem=TensorProblem(), model=FeedForward(2, 1), use_lt=False
|
||||
)
|
||||
supervised_solver_lt = SupervisedSolver(
|
||||
problem=LabelTensorProblem(), model=FeedForward(2, 1), use_lt=True
|
||||
)
|
||||
|
||||
poisson_problem = Poisson()
|
||||
poisson_problem.conditions["data"] = Condition(
|
||||
input=LabelTensor(torch.rand(20, 2) * 10, ["x", "y"]),
|
||||
target=LabelTensor(torch.rand(20, 1) * 10, ["u"]),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scale_fn", [torch.std, torch.var])
|
||||
@pytest.mark.parametrize("shift_fn", [torch.mean, torch.median])
|
||||
@pytest.mark.parametrize("apply_to", ["input", "target"])
|
||||
@pytest.mark.parametrize("stage", ["train", "validate", "test", "all"])
|
||||
def test_init(scale_fn, shift_fn, apply_to, stage):
|
||||
normalizer = NormalizerDataCallback(
|
||||
scale_fn=scale_fn, shift_fn=shift_fn, apply_to=apply_to, stage=stage
|
||||
)
|
||||
assert normalizer.scale_fn == scale_fn
|
||||
assert normalizer.shift_fn == shift_fn
|
||||
assert normalizer.apply_to == apply_to
|
||||
assert normalizer.stage == stage
|
||||
|
||||
|
||||
def test_init_invalid_scale():
|
||||
with pytest.raises(ValueError):
|
||||
NormalizerDataCallback(scale_fn=1)
|
||||
|
||||
|
||||
def test_init_invalid_shift():
|
||||
with pytest.raises(ValueError):
|
||||
NormalizerDataCallback(shift_fn=1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid_apply_to", ["inputt", "targett", 1])
|
||||
def test_init_invalid_apply_to(invalid_apply_to):
|
||||
with pytest.raises(ValueError):
|
||||
NormalizerDataCallback(apply_to=invalid_apply_to)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid_stage", ["trainn", "validatee", 1])
|
||||
def test_init_invalid_stage(invalid_stage):
|
||||
with pytest.raises(ValueError):
|
||||
NormalizerDataCallback(stage=invalid_stage)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"solver", [supervised_solver_lt, supervised_solver_no_lt]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"fn", [[torch.std, torch.mean], [torch.var, torch.median]]
|
||||
)
|
||||
@pytest.mark.parametrize("apply_to", ["input", "target"])
|
||||
@pytest.mark.parametrize("stage", ["all", "train", "validate", "test"])
|
||||
def test_setup(solver, fn, stage, apply_to):
|
||||
scale_fn, shift_fn = fn
|
||||
trainer = Trainer(
|
||||
solver=solver,
|
||||
callbacks=NormalizerDataCallback(
|
||||
scale_fn=scale_fn, shift_fn=shift_fn, stage=stage, apply_to=apply_to
|
||||
),
|
||||
max_epochs=1,
|
||||
train_size=0.4,
|
||||
val_size=0.3,
|
||||
test_size=0.3,
|
||||
shuffle=False,
|
||||
)
|
||||
trainer_copy = deepcopy(trainer)
|
||||
trainer_copy.data_module.setup("fit")
|
||||
trainer_copy.data_module.setup("test")
|
||||
trainer.train()
|
||||
trainer.test()
|
||||
|
||||
normalizer = trainer.callbacks[0].normalizer
|
||||
|
||||
for cond in ["data1", "data2"]:
|
||||
scale = scale_fn(
|
||||
trainer_copy.data_module.train_dataset.conditions_dict[cond][
|
||||
apply_to
|
||||
]
|
||||
)
|
||||
shift = shift_fn(
|
||||
trainer_copy.data_module.train_dataset.conditions_dict[cond][
|
||||
apply_to
|
||||
]
|
||||
)
|
||||
assert "scale" in normalizer[cond]
|
||||
assert "shift" in normalizer[cond]
|
||||
assert normalizer[cond]["scale"] - scale < 1e-5
|
||||
assert normalizer[cond]["shift"] - shift < 1e-5
|
||||
for ds_name in stage_map[stage]:
|
||||
dataset = getattr(trainer.data_module, ds_name, None)
|
||||
old_dataset = getattr(trainer_copy.data_module, ds_name, None)
|
||||
current_points = dataset.conditions_dict[cond][apply_to]
|
||||
old_points = old_dataset.conditions_dict[cond][apply_to]
|
||||
expected = (old_points - shift) / scale
|
||||
assert torch.allclose(current_points, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fn", [[torch.std, torch.mean], [torch.var, torch.median]]
|
||||
)
|
||||
@pytest.mark.parametrize("apply_to", ["input"])
|
||||
@pytest.mark.parametrize("stage", ["all", "train", "validate", "test"])
|
||||
def test_setup_pinn(fn, stage, apply_to):
|
||||
scale_fn, shift_fn = fn
|
||||
pinn = PINN(
|
||||
problem=poisson_problem,
|
||||
model=FeedForward(2, 1),
|
||||
)
|
||||
poisson_problem.discretise_domain(n=10)
|
||||
trainer = Trainer(
|
||||
solver=pinn,
|
||||
callbacks=NormalizerDataCallback(
|
||||
scale_fn=scale_fn,
|
||||
shift_fn=shift_fn,
|
||||
stage=stage,
|
||||
apply_to=apply_to,
|
||||
),
|
||||
max_epochs=1,
|
||||
train_size=0.4,
|
||||
val_size=0.3,
|
||||
test_size=0.3,
|
||||
shuffle=False,
|
||||
)
|
||||
|
||||
trainer_copy = deepcopy(trainer)
|
||||
trainer_copy.data_module.setup("fit")
|
||||
trainer_copy.data_module.setup("test")
|
||||
trainer.train()
|
||||
trainer.test()
|
||||
|
||||
conditions = trainer.callbacks[0].normalizer.keys()
|
||||
assert "data" in conditions
|
||||
assert len(conditions) == 1
|
||||
normalizer = trainer.callbacks[0].normalizer
|
||||
cond = "data"
|
||||
|
||||
scale = scale_fn(
|
||||
trainer_copy.data_module.train_dataset.conditions_dict[cond][apply_to]
|
||||
)
|
||||
shift = shift_fn(
|
||||
trainer_copy.data_module.train_dataset.conditions_dict[cond][apply_to]
|
||||
)
|
||||
assert "scale" in normalizer[cond]
|
||||
assert "shift" in normalizer[cond]
|
||||
assert normalizer[cond]["scale"] - scale < 1e-5
|
||||
assert normalizer[cond]["shift"] - shift < 1e-5
|
||||
for ds_name in stage_map[stage]:
|
||||
dataset = getattr(trainer.data_module, ds_name, None)
|
||||
old_dataset = getattr(trainer_copy.data_module, ds_name, None)
|
||||
current_points = dataset.conditions_dict[cond][apply_to]
|
||||
old_points = old_dataset.conditions_dict[cond][apply_to]
|
||||
expected = (old_points - shift) / scale
|
||||
assert torch.allclose(current_points, expected)
|
||||
|
||||
|
||||
def test_setup_graph_dataset():
|
||||
solver = SupervisedSolver(
|
||||
problem=GraphProblem(), model=FeedForward(2, 1), use_lt=False
|
||||
)
|
||||
trainer = Trainer(
|
||||
solver=solver,
|
||||
callbacks=NormalizerDataCallback(
|
||||
scale_fn=torch.std,
|
||||
shift_fn=torch.mean,
|
||||
stage="all",
|
||||
apply_to="input",
|
||||
),
|
||||
max_epochs=1,
|
||||
train_size=0.4,
|
||||
val_size=0.3,
|
||||
test_size=0.3,
|
||||
shuffle=False,
|
||||
)
|
||||
with pytest.raises(NotImplementedError):
|
||||
trainer.train()
|
||||
Reference in New Issue
Block a user