From f57a08b87589fdb58efbf3103ace83c075d16c5b Mon Sep 17 00:00:00 2001 From: Nicola Demo Date: Wed, 28 Jun 2023 15:13:47 +0200 Subject: [PATCH] device for sample points in absProblem (#132) * device for sample points in absProblem --- pina/problem/abstract_problem.py | 6 +++++- tests/test_solvers/test_pinn.py | 14 ++++++++++++-- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py index e67b15f..c3abf4b 100644 --- a/pina/problem/abstract_problem.py +++ b/pina/problem/abstract_problem.py @@ -111,7 +111,7 @@ class AbstractProblem(metaclass=ABCMeta): continue self.input_pts[condition_name] = samples - def discretise_domain(self, n, mode = 'random', variables = 'all', locations = 'all'): + def discretise_domain(self, n, mode = 'random', variables = 'all', locations = 'all', device=None): """ Generate a set of points to span the `Location` of all the conditions of the problem. @@ -196,6 +196,10 @@ class AbstractProblem(metaclass=ABCMeta): ] + already_sampled pts = merge_tensors(samples) self.input_pts[location] = pts + + if device: + self.input_pts[location] = self.input_pts[location].to(device=device) #TODO better fix + # setting the grad self.input_pts[location].requires_grad_(True) self.input_pts[location].retain_grad() diff --git a/tests/test_solvers/test_pinn.py b/tests/test_solvers/test_pinn.py index 4e533f6..dd9d0d8 100644 --- a/tests/test_solvers/test_pinn.py +++ b/tests/test_solvers/test_pinn.py @@ -105,8 +105,18 @@ def test_train_extra_feats_cpu(): pinn = PINN(problem = poisson_problem, model=model_extra_feats, extra_features=extra_feats) trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'}) trainer.train() - """ +def test_train_gpu(): #TODO fix ASAP + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + poisson_problem.discretise_domain(n, 'grid', locations=['D']) + poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu + pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) + trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'}) + trainer.train() + def test_train_2(): boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] n = 10 @@ -212,4 +222,4 @@ if torch.cuda.is_available(): pinn.discretise_domain(n, 'grid', locations=boundaries) pinn.discretise_domain(n, 'grid', locations=['D']) pinn.train(5) -""" \ No newline at end of file +"""