46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
from pina.callbacks import SwitchOptimizer
|
|
import torch
|
|
import pytest
|
|
|
|
from pina.solvers import PINN
|
|
from pina.trainer import Trainer
|
|
from pina.model import FeedForward
|
|
from pina.problem.zoo import Poisson2DSquareProblem as Poisson
|
|
|
|
# make the problem
|
|
poisson_problem = Poisson()
|
|
boundaries = ['nil_g1', 'nil_g2', 'nil_g3', 'nil_g4']
|
|
n = 10
|
|
poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
|
|
model = FeedForward(len(poisson_problem.input_variables),
|
|
len(poisson_problem.output_variables))
|
|
|
|
# make the solver
|
|
solver = PINN(problem=poisson_problem, model=model)
|
|
|
|
|
|
def test_switch_optimizer_constructor():
|
|
SwitchOptimizer(new_optimizers=torch.optim.Adam,
|
|
new_optimizers_kwargs={'lr': 0.01},
|
|
epoch_switch=10)
|
|
|
|
with pytest.raises(ValueError):
|
|
SwitchOptimizer(new_optimizers=[torch.optim.Adam, torch.optim.Adam],
|
|
new_optimizers_kwargs=[{
|
|
'lr': 0.01
|
|
}],
|
|
epoch_switch=10)
|
|
|
|
|
|
def test_switch_optimizer_routine():
|
|
# make the trainer
|
|
trainer = Trainer(solver=solver,
|
|
callbacks=[
|
|
SwitchOptimizer(new_optimizers=torch.optim.LBFGS,
|
|
new_optimizers_kwargs={'lr': 0.01},
|
|
epoch_switch=3)
|
|
],
|
|
accelerator='cpu',
|
|
max_epochs=5)
|
|
trainer.train()
|