9c9d4fe7e4
* Fix bug in Collector with Graph data * Add comments in DataModule class and bug fix in collate
256 lines
8.8 KiB
Python
256 lines
8.8 KiB
Python
import torch
|
|
import pytest
|
|
from pina.data import PinaDataModule
|
|
from pina.data.dataset import PinaTensorDataset, PinaGraphDataset
|
|
from pina.problem.zoo import SupervisedProblem
|
|
from pina.graph import RadiusGraph
|
|
from pina.data.data_module import DummyDataloader
|
|
from pina import Trainer
|
|
from pina.solvers import SupervisedSolver
|
|
from torch_geometric.data import Batch
|
|
from torch.utils.data import DataLoader
|
|
|
|
input_tensor = torch.rand((100, 10))
|
|
output_tensor = torch.rand((100, 2))
|
|
|
|
x = torch.rand((100, 50, 10))
|
|
pos = torch.rand((100, 50, 2))
|
|
input_graph = RadiusGraph(x, pos, r=.1, build_edge_attr=True)
|
|
output_graph = torch.rand((100, 50, 10))
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_, output_",
|
|
[
|
|
(input_tensor, output_tensor),
|
|
(input_graph, output_graph)
|
|
]
|
|
)
|
|
def test_constructor(input_, output_):
|
|
problem = SupervisedProblem(input_=input_, output_=output_)
|
|
PinaDataModule(problem)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_, output_",
|
|
[
|
|
(input_tensor, output_tensor),
|
|
(input_graph, output_graph)
|
|
]
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"train_size, val_size, test_size",
|
|
[
|
|
(.7, .2, .1),
|
|
(.7, .3, 0)
|
|
]
|
|
)
|
|
def test_setup_train(input_, output_, train_size, val_size, test_size):
|
|
problem = SupervisedProblem(input_=input_, output_=output_)
|
|
dm = PinaDataModule(problem, train_size=train_size,
|
|
val_size=val_size, test_size=test_size)
|
|
dm.setup()
|
|
assert hasattr(dm, "train_dataset")
|
|
if isinstance(input_, torch.Tensor):
|
|
assert isinstance(dm.train_dataset, PinaTensorDataset)
|
|
else:
|
|
assert isinstance(dm.train_dataset, PinaGraphDataset)
|
|
# assert len(dm.train_dataset) == int(len(input_) * train_size)
|
|
if test_size > 0:
|
|
assert hasattr(dm, "test_dataset")
|
|
assert dm.test_dataset is None
|
|
else:
|
|
assert not hasattr(dm, "test_dataset")
|
|
assert hasattr(dm, "val_dataset")
|
|
if isinstance(input_, torch.Tensor):
|
|
assert isinstance(dm.val_dataset, PinaTensorDataset)
|
|
else:
|
|
assert isinstance(dm.val_dataset, PinaGraphDataset)
|
|
# assert len(dm.val_dataset) == int(len(input_) * val_size)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_, output_",
|
|
[
|
|
(input_tensor, output_tensor),
|
|
(input_graph, output_graph)
|
|
]
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"train_size, val_size, test_size",
|
|
[
|
|
(.7, .2, .1),
|
|
(0., 0., 1.)
|
|
]
|
|
)
|
|
def test_setup_test(input_, output_, train_size, val_size, test_size):
|
|
problem = SupervisedProblem(input_=input_, output_=output_)
|
|
dm = PinaDataModule(problem, train_size=train_size,
|
|
val_size=val_size, test_size=test_size)
|
|
dm.setup(stage='test')
|
|
if train_size > 0:
|
|
assert hasattr(dm, "train_dataset")
|
|
assert dm.train_dataset is None
|
|
else:
|
|
assert not hasattr(dm, "train_dataset")
|
|
if val_size > 0:
|
|
assert hasattr(dm, "val_dataset")
|
|
assert dm.val_dataset is None
|
|
else:
|
|
assert not hasattr(dm, "val_dataset")
|
|
|
|
assert hasattr(dm, "test_dataset")
|
|
if isinstance(input_, torch.Tensor):
|
|
assert isinstance(dm.test_dataset, PinaTensorDataset)
|
|
else:
|
|
assert isinstance(dm.test_dataset, PinaGraphDataset)
|
|
# assert len(dm.test_dataset) == int(len(input_) * test_size)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_, output_",
|
|
[
|
|
(input_tensor, output_tensor),
|
|
(input_graph, output_graph)
|
|
]
|
|
)
|
|
def test_dummy_dataloader(input_, output_):
|
|
problem = SupervisedProblem(input_=input_, output_=output_)
|
|
solver = SupervisedSolver(problem=problem, model=torch.nn.Linear(10, 10))
|
|
trainer = Trainer(solver, batch_size=None, train_size=.7,
|
|
val_size=.3, test_size=0.)
|
|
dm = trainer.data_module
|
|
dm.setup()
|
|
dm.trainer = trainer
|
|
dataloader = dm.train_dataloader()
|
|
assert isinstance(dataloader, DummyDataloader)
|
|
assert len(dataloader) == 1
|
|
data = next(dataloader)
|
|
assert isinstance(data, list)
|
|
assert isinstance(data[0], tuple)
|
|
if isinstance(input_, RadiusGraph):
|
|
assert isinstance(data[0][1]['input_points'], Batch)
|
|
else:
|
|
assert isinstance(data[0][1]['input_points'], torch.Tensor)
|
|
assert isinstance(data[0][1]['output_points'], torch.Tensor)
|
|
|
|
dataloader = dm.val_dataloader()
|
|
assert isinstance(dataloader, DummyDataloader)
|
|
assert len(dataloader) == 1
|
|
data = next(dataloader)
|
|
assert isinstance(data, list)
|
|
assert isinstance(data[0], tuple)
|
|
if isinstance(input_, RadiusGraph):
|
|
assert isinstance(data[0][1]['input_points'], Batch)
|
|
else:
|
|
assert isinstance(data[0][1]['input_points'], torch.Tensor)
|
|
assert isinstance(data[0][1]['output_points'], torch.Tensor)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_, output_",
|
|
[
|
|
(input_tensor, output_tensor),
|
|
(input_graph, output_graph)
|
|
]
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"automatic_batching",
|
|
[
|
|
True, False
|
|
]
|
|
)
|
|
def test_dataloader(input_, output_, automatic_batching):
|
|
problem = SupervisedProblem(input_=input_, output_=output_)
|
|
solver = SupervisedSolver(problem=problem, model=torch.nn.Linear(10, 10))
|
|
trainer = Trainer(solver, batch_size=10, train_size=.7, val_size=.3,
|
|
test_size=0., automatic_batching=automatic_batching)
|
|
dm = trainer.data_module
|
|
dm.setup()
|
|
dm.trainer = trainer
|
|
dataloader = dm.train_dataloader()
|
|
assert isinstance(dataloader, DataLoader)
|
|
assert len(dataloader) == 7
|
|
data = next(iter(dataloader))
|
|
assert isinstance(data, dict)
|
|
if isinstance(input_, RadiusGraph):
|
|
assert isinstance(data['data']['input_points'], Batch)
|
|
else:
|
|
assert isinstance(data['data']['input_points'], torch.Tensor)
|
|
assert isinstance(data['data']['output_points'], torch.Tensor)
|
|
|
|
dataloader = dm.val_dataloader()
|
|
assert isinstance(dataloader, DataLoader)
|
|
assert len(dataloader) == 3
|
|
data = next(iter(dataloader))
|
|
assert isinstance(data, dict)
|
|
if isinstance(input_, RadiusGraph):
|
|
assert isinstance(data['data']['input_points'], Batch)
|
|
else:
|
|
assert isinstance(data['data']['input_points'], torch.Tensor)
|
|
assert isinstance(data['data']['output_points'], torch.Tensor)
|
|
|
|
from pina import LabelTensor
|
|
|
|
input_tensor = LabelTensor(torch.rand((100, 3)), ['u', 'v', 'w'])
|
|
output_tensor = LabelTensor(torch.rand((100, 3)), ['u', 'v', 'w'])
|
|
|
|
x = LabelTensor(torch.rand((100, 50, 3)), ['u', 'v', 'w'])
|
|
pos = LabelTensor(torch.rand((100, 50, 2)), ['x', 'y'])
|
|
input_graph = RadiusGraph(x, pos, r=.1, build_edge_attr=True)
|
|
output_graph = LabelTensor(torch.rand((100, 50, 3)), ['u', 'v', 'w'])
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_, output_",
|
|
[
|
|
(input_tensor, output_tensor),
|
|
(input_graph, output_graph)
|
|
]
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"automatic_batching",
|
|
[
|
|
True, False
|
|
]
|
|
)
|
|
def test_dataloader_labels(input_, output_, automatic_batching):
|
|
problem = SupervisedProblem(input_=input_, output_=output_)
|
|
solver = SupervisedSolver(problem=problem, model=torch.nn.Linear(10, 10))
|
|
trainer = Trainer(solver, batch_size=10, train_size=.7, val_size=.3,
|
|
test_size=0., automatic_batching=automatic_batching)
|
|
dm = trainer.data_module
|
|
dm.setup()
|
|
dm.trainer = trainer
|
|
dataloader = dm.train_dataloader()
|
|
assert isinstance(dataloader, DataLoader)
|
|
assert len(dataloader) == 7
|
|
data = next(iter(dataloader))
|
|
assert isinstance(data, dict)
|
|
if isinstance(input_, RadiusGraph):
|
|
assert isinstance(data['data']['input_points'], Batch)
|
|
assert isinstance(data['data']['input_points'].x, LabelTensor)
|
|
assert data['data']['input_points'].x.labels == ['u', 'v', 'w']
|
|
assert data['data']['input_points'].pos.labels == ['x', 'y']
|
|
else:
|
|
assert isinstance(data['data']['input_points'], LabelTensor)
|
|
assert data['data']['input_points'].labels == ['u', 'v', 'w']
|
|
assert isinstance(data['data']['output_points'], LabelTensor)
|
|
assert data['data']['output_points'].labels == ['u', 'v', 'w']
|
|
|
|
dataloader = dm.val_dataloader()
|
|
assert isinstance(dataloader, DataLoader)
|
|
assert len(dataloader) == 3
|
|
data = next(iter(dataloader))
|
|
assert isinstance(data, dict)
|
|
if isinstance(input_, RadiusGraph):
|
|
assert isinstance(data['data']['input_points'], Batch)
|
|
assert isinstance(data['data']['input_points'].x, LabelTensor)
|
|
assert data['data']['input_points'].x.labels == ['u', 'v', 'w']
|
|
assert data['data']['input_points'].pos.labels == ['x', 'y']
|
|
else:
|
|
assert isinstance(data['data']['input_points'], torch.Tensor)
|
|
assert isinstance(data['data']['input_points'], LabelTensor)
|
|
assert data['data']['input_points'].labels == ['u', 'v', 'w']
|
|
assert isinstance(data['data']['output_points'], torch.Tensor)
|
|
assert data['data']['output_points'].labels == ['u', 'v', 'w']
|
|
test_dataloader_labels(input_graph, output_graph, True) |