Rename files
This commit is contained in:
committed by
Nicola Demo
parent
886bd23fdb
commit
ff43a7492b
256
tests/test_data/test_data_module.py
Normal file
256
tests/test_data/test_data_module.py
Normal file
@@ -0,0 +1,256 @@
|
||||
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.solver 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)
|
||||
Reference in New Issue
Block a user