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

batch_enhancement (#51)

Browse Source
This commit is contained in:
Dario Coscia
2022-12-12 11:09:20 +01:00
committed by GitHub
parent d70f5e730a
commit dbd78c9cf3
4 changed files with 236 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
pina/label_tensor.py
View File

@@ -106,6 +106,14 @@ class LabelTensor(torch.Tensor):
new.data = tmp.data
return new
def select(self, *args, **kwargs):
"""
Performs Tensor selection. For more details, see :meth:`torch.Tensor.select`.
"""
tmp = super().select(*args, **kwargs)
tmp._labels = self._labels
return tmp
def extract(self, label_to_extract):
"""
Extract the subset of the original tensor by returning all the columns

99
pina/pinn.py
View File

@@ -3,7 +3,8 @@ import torch
from .problem import AbstractProblem
from .label_tensor import LabelTensor
from .utils import merge_tensors
from .utils import merge_tensors, PinaDataset
torch.pi = torch.acos(torch.zeros(1)).item() * 2 # which is 3.1415927410125732
@@ -16,17 +17,27 @@ class PINN(object):
optimizer=torch.optim.Adam,
lr=0.001,
regularizer=0.00001,
batch_size=None,
dtype=torch.float32,
device='cpu',
error_norm='mse'):
'''
:param Problem problem: the formualation of the problem.
:param torch.nn.Module model: the neural network model to use.
:param torch.optim optimizer: the neural network optimizer to use;
default is `torch.optim.Adam`.
:param float lr: the learning rate; default is 0.001.
:param float regularizer: the coefficient for L2 regularizer term.
:param type dtype: the data type to use for the model. Valid option are
`torch.float32` and `torch.float64` (`torch.float16` only on GPU);
default is `torch.float64`.
:param string device: the device used for training; default 'cpu'
option include 'cuda' if cuda is available.
:param string/int error_norm: the loss function used as minimizer,
default mean square error 'mse'. If string options include mean
error 'me' and mean square error 'mse'. If int, the p-norm is
calculated where p is specifined by the int input.
:param int batch_size: batch size for the dataloader; default 5.
'''
if dtype == torch.float64:
@@ -59,6 +70,9 @@ class PINN(object):
self.optimizer = optimizer(
self.model.parameters(), lr=lr, weight_decay=regularizer)
self.batch_size = batch_size
self.data_set = PinaDataset(self)
@property
def problem(self):
return self._problem
@@ -110,7 +124,6 @@ class PINN(object):
checkpoint = torch.load(filename)
self.model.load_state_dict(checkpoint['model_state'])
self.optimizer = checkpoint['optimizer_class'](self.model.parameters())
self.optimizer.load_state_dict(checkpoint['optimizer_state'])
@@ -121,6 +134,39 @@ class PINN(object):
return self
def _create_dataloader(self):
"""Private method for creating dataloader
:return: dataloader
:rtype: torch.utils.data.DataLoader
"""
if self.batch_size is None:
return [self.input_pts]
def custom_collate(batch):
# extracting pts labels
_, pts = list(batch[0].items())[0]
labels = pts.labels
# calling default torch collate
collate_res = default_collate(batch)
# save collate result in dict
res = {}
for key, val in collate_res.items():
val.labels = labels
res[key] = val
return res
# creating dataset, list of dataset for each location
datasets = [MyDataSet(key, val)
for key, val in self.input_pts.items()]
# creating dataloader
dataloaders = [DataLoader(dataset=dat,
batch_size=self.batch_size,
collate_fn=custom_collate)
for dat in datasets]
return dict(zip(self.input_pts.keys(), dataloaders))
def span_pts(self, *args, **kwargs):
"""
>>> pinn.span_pts(n=10, mode='grid')
@@ -160,55 +206,65 @@ class PINN(object):
# TODO
# pts = pts.double()
pts = pts.to(dtype=self.dtype, device=self.device)
pts.requires_grad_(True)
pts.retain_grad()
self.input_pts[location] = pts
def train(self, stop=100, frequency_print=2, save_loss=1, trial=None):
epoch = 0
data_loader = self.data_set.dataloader
header = []
for condition_name in self.problem.conditions:
condition = self.problem.conditions[condition_name]
if (hasattr(condition, 'function') and
isinstance(condition.function, list)):
if hasattr(condition, 'function'):
if isinstance(condition.function, list):
for function in condition.function:
header.append(f'{condition_name}{function.__name__}')
else:
continue
header.append(f'{condition_name}')
while True:
losses = []
for condition_name in self.problem.conditions:
condition = self.problem.conditions[condition_name]
for batch in data_loader[condition_name]:
single_loss = []
if hasattr(condition, 'function'):
pts = self.input_pts[condition_name]
pts = batch[condition_name]
pts = pts.to(dtype=self.dtype, device=self.device)
pts.requires_grad_(True)
pts.retain_grad()
predicted = self.model(pts)
for function in condition.function:
residuals = function(pts, predicted)
local_loss = (
condition.data_weight*self._compute_norm(
residuals))
losses.append(local_loss)
single_loss.append(local_loss)
elif hasattr(condition, 'output_points'):
pts = condition.input_points
pts = condition.input_points.to(
dtype=self.dtype, device=self.device)
predicted = self.model(pts)
residuals = predicted - condition.output_points
local_loss = (
condition.data_weight*self._compute_norm(residuals))
losses.append(local_loss)
single_loss.append(local_loss)
self.optimizer.zero_grad()
sum(losses).backward()
sum(single_loss).backward()
self.optimizer.step()
losses.append(sum(single_loss))
if save_loss and (epoch % save_loss == 0 or epoch == 0):
self.history_loss[epoch] = [
loss.detach().item() for loss in losses]
@@ -221,7 +277,8 @@ class PINN(object):
if isinstance(stop, int):
if epoch == stop:
print('[epoch {:05d}] {:.6e} '.format(self.trained_epoch, sum(losses).item()), end='')
print('[epoch {:05d}] {:.6e} '.format(
self.trained_epoch, sum(losses).item()), end='')
for loss in losses:
print('{:.6e} '.format(loss.item()), end='')
print()
@@ -236,7 +293,8 @@ class PINN(object):
print('{:12.12s} '.format(name), end='')
print()
print('[epoch {:05d}] {:.6e} '.format(self.trained_epoch, sum(losses).item()), end='')
print('[epoch {:05d}] {:.6e} '.format(
self.trained_epoch, sum(losses).item()), end='')
for loss in losses:
print('{:.6e} '.format(loss.item()), end='')
print()
@@ -246,7 +304,6 @@ class PINN(object):
return sum(losses).item()
def error(self, dtype='l2', res=100):
import numpy as np
@@ -261,7 +318,8 @@ class PINN(object):
grids_container = self.problem.data_solution['grid']
Z_true = self.problem.data_solution['grid_solution']
try:
unrolled_pts = torch.tensor([t.flatten() for t in grids_container]).T.to(dtype=self.dtype, device=self.device)
unrolled_pts = torch.tensor([t.flatten() for t in grids_container]).T.to(
dtype=self.dtype, device=self.device)
Z_pred = self.model(unrolled_pts)
Z_pred = Z_pred.detach().numpy().reshape(grids_container[0].shape)
@@ -273,4 +331,5 @@ class PINN(object):
except:
print("")
print("Something went wrong...")
print("Not able to compute the error. Please pass a data solution or a true solution")
print(
"Not able to compute the error. Please pass a data solution or a true solution")

64
pina/utils.py
View File

@@ -1,5 +1,7 @@
"""Utils module"""
from functools import reduce
import torch
from torch.utils.data import DataLoader, default_collate, ConcatDataset
from .label_tensor import LabelTensor
@@ -45,3 +47,65 @@ def merge_two_tensors(tensor1, tensor2):
tensor2 = LabelTensor(tensor2.repeat_interleave(n1, dim=0),
labels=tensor2.labels)
return tensor1.append(tensor2)
class PinaDataset():
def __init__(self, pinn) -> None:
self.pinn = pinn
@property
def dataloader(self):
return self._create_dataloader()
@property
def dataset(self):
return [self.SampleDataset(key, val)
for key, val in self.input_pts.items()]
def _create_dataloader(self):
"""Private method for creating dataloader
:return: dataloader
:rtype: torch.utils.data.DataLoader
"""
if self.pinn.batch_size is None:
return {key: [{key: val}] for key, val in self.pinn.input_pts.items()}
def custom_collate(batch):
# extracting pts labels
_, pts = list(batch[0].items())[0]
labels = pts.labels
# calling default torch collate
collate_res = default_collate(batch)
# save collate result in dict
res = {}
for key, val in collate_res.items():
val.labels = labels
res[key] = val
return res
# creating dataset, list of dataset for each location
datasets = [self.SampleDataset(key, val)
for key, val in self.pinn.input_pts.items()]
# creating dataloader
dataloaders = [DataLoader(dataset=dat,
batch_size=self.pinn.batch_size,
collate_fn=custom_collate)
for dat in datasets]
return dict(zip(self.pinn.input_pts.keys(), dataloaders))
class SampleDataset(torch.utils.data.Dataset):
def __init__(self, location, tensor):
self._tensor = tensor
self._location = location
self._len = len(tensor)
def __getitem__(self, index):
tensor = self._tensor.select(0, index)
return {self._location: tensor}
def __len__(self):
return self._len

46
tests/test_pinn.py
View File

@@ -37,13 +37,16 @@ class Poisson(SpatialProblem):
truth_solution = poisson_sol
problem = Poisson()
model = FeedForward(problem.input_variables, problem.output_variables)
def test_constructor():
PINN(problem, model)
def test_span_pts():
pinn = PINN(problem, model)
n = 10
@@ -60,6 +63,7 @@ def test_span_pts():
pinn.span_pts(n, 'random', locations=['D'])
assert pinn.input_pts['D'].shape[0] == n
def test_train():
pinn = PINN(problem, model)
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
@@ -68,6 +72,7 @@ def test_train():
pinn.span_pts(n, 'grid', locations=['D'])
pinn.train(5)
def test_train():
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
n = 10
@@ -79,3 +84,44 @@ def test_train():
pinn.span_pts(n, 'grid', locations=['D'])
pinn.train(50, save_loss=i)
assert list(pinn.history_loss.keys()) == truth_key
def test_train_batch():
pinn = PINN(problem, model, batch_size=6)
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
n = 10
pinn.span_pts(n, 'grid', boundaries)
pinn.span_pts(n, 'grid', locations=['D'])
pinn.train(5)
def test_train_batch():
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
n = 10
expected_keys = [[], list(range(0, 50, 3))]
param = [0, 3]
for i, truth_key in zip(param, expected_keys):
pinn = PINN(problem, model, batch_size=6)
pinn.span_pts(n, 'grid', boundaries)
pinn.span_pts(n, 'grid', locations=['D'])
pinn.train(50, save_loss=i)
assert list(pinn.history_loss.keys()) == truth_key
if torch.cuda.is_available():
def test_gpu_train():
pinn = PINN(problem, model, batch_size=20, device='cuda')
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
n = 100
pinn.span_pts(n, 'grid', boundaries)
pinn.span_pts(n, 'grid', locations=['D'])
pinn.train(5)
def test_gpu_train_nobatch():
pinn = PINN(problem, model, batch_size=None, device='cuda')
boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
n = 100
pinn.span_pts(n, 'grid', boundaries)
pinn.span_pts(n, 'grid', locations=['D'])
pinn.train(5)