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

CPU/GPU/TPU training (#159)

Browse Source 
* device training

---------

Co-authored-by: Dario Coscia <dcoscia@lovelace.maths.sissa.it>
Co-authored-by: Dario Coscia <dariocoscia@Dario-Coscia.local>
This commit is contained in:
Dario Coscia
2023-07-19 17:19:08 +02:00
committed by Nicola Demo
parent 38ecebd44b
commit 92e0e4920b
4 changed files with 62 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
pina/dataset.py
View File

@@ -1,5 +1,6 @@
""" """
from torch.utils.data import Dataset, DataLoader
import functools
class PinaDataset():
@@ -48,7 +49,30 @@ class PinaDataset():
# TODO: working also for datapoints
class DummyLoader:
def __init__(self, data) -> None:
def __init__(self, data, device) -> None:
# TODO: We need to make a dataset somehow
# and the PINADataset needs to have a method
# to send points to device
# now we simply do it here
# send data to device
def convert_tensors(pts, device):
pts = pts.to(device)
pts.requires_grad_(True)
pts.retain_grad()
return pts
for location, pts in data.items():
if isinstance(pts, (tuple, list)):
pts = tuple(map(functools.partial(convert_tensors, device=device),pts))
else:
pts = pts.to(device)
pts = pts.requires_grad_(True)
pts.retain_grad()
data[location] = pts
# iterator
self.data = [data]
def __iter__(self):

14
pina/problem/abstract_problem.py
View File

@@ -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', device=None):
def discretise_domain(self, n, mode = 'random', variables = 'all', locations = 'all'):
"""
Generate a set of points to span the `Location` of all the conditions of
the problem.
@@ -129,9 +129,9 @@ class AbstractProblem(metaclass=ABCMeta):
:type locations: str, optional
:Example:
>>> pinn.span_pts(n=10, mode='grid')
>>> pinn.span_pts(n=10, mode='grid', location=['bound1'])
>>> pinn.span_pts(n=10, mode='grid', variables=['x'])
>>> pinn.discretise_domain(n=10, mode='grid')
>>> pinn.discretise_domain(n=10, mode='grid', location=['bound1'])
>>> pinn.discretise_domain(n=10, mode='grid', variables=['x'])
.. warning::
``random`` is currently the only implemented ``mode`` for all geometries, i.e.
@@ -200,12 +200,6 @@ class AbstractProblem(metaclass=ABCMeta):
# the condition is sampled if input_pts contains all labels
if sorted(self.input_pts[location].labels) == sorted(self.input_variables):
self._have_sampled_points[location] = True
# setting device
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()
@property
def have_sampled_points(self):

5
pina/trainer.py
View File

@@ -10,6 +10,9 @@ class Trainer(pl.Trainer):
def __init__(self, solver, kwargs={}):
super().__init__(**kwargs)
# get accellerator
device = self._accelerator_connector._accelerator_flag
# check inheritance consistency for solver
check_consistency(solver, SolverInterface)
self._model = solver
@@ -23,7 +26,7 @@ class Trainer(pl.Trainer):
'in the provided locations.')
# TODO: make a better dataloader for train
self._loader = DummyLoader(solver.problem.input_pts)
self._loader = DummyLoader(solver.problem.input_pts, device)
def train(self): # TODO add kwargs and lightining capabilities

45
tests/test_solvers/test_pinn.py
View File

@@ -20,8 +20,8 @@ def laplace_equation(input_, output_):
return nabla_u - force_term
my_laplace = Equation(laplace_equation)
in_ = LabelTensor(torch.tensor([[0., 1.]], requires_grad=True), ['x', 'y'])
out_ = LabelTensor(torch.tensor([[0.]], requires_grad=True), ['u'])
in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
class Poisson(SpatialProblem):
output_variables = ['u']
@@ -41,7 +41,7 @@ class Poisson(SpatialProblem):
location=CartesianDomain({'x': 0, 'y': [0, 1]}),
equation=FixedValue(0.0)),
'D': Condition(
location=CartesianDomain({'x': [0, 1], 'y': [0, 1]}),
input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
equation=my_laplace),
'data': Condition(
input_points=in_,
@@ -91,21 +91,25 @@ def test_train_cpu():
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
n = 10
poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
poisson_problem.discretise_domain(n, 'grid', locations=['D'])
pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
trainer.train()
def test_train_cpu_sampling_few_vars():
poisson_problem = Poisson()
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
n = 10
poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
poisson_problem.discretise_domain(n, 'random', locations=['D'], variables=['x'])
poisson_problem.discretise_domain(n, 'random', locations=['D'], variables=['y'])
pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
trainer.train()
# # TODO fix asap. Basically sampling few variables
# # works only if both variables are in a range.
# # if one is fixed and the other not, this will
# # not work. This test also needs to be fixed and
# # insert in test problem not in test pinn.
# def test_train_cpu_sampling_few_vars():
# poisson_problem = Poisson()
# boundaries = ['gamma1', 'gamma2', 'gamma3']
# n = 10
# poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
# poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x'])
# poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y'])
# pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
# trainer.train()
def test_train_extra_feats_cpu():
@@ -113,17 +117,26 @@ def test_train_extra_feats_cpu():
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
n = 10
poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
poisson_problem.discretise_domain(n, 'grid', locations=['D'])
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()
# TODO, fix GitHub actions to run also on GPU
# def test_train_gpu():
# poisson_problem = Poisson()
# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
# n = 10
# poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
# 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_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'})