ab6ca78d85
* Simplifying Graph class and adjust tests --------- Co-authored-by: Dario Coscia <dariocos99@gmail.com>
117 lines
3.5 KiB
Python
117 lines
3.5 KiB
Python
import pytest
|
|
import torch
|
|
from pina.graph import KNNGraph
|
|
from pina.model import GraphNeuralOperator
|
|
from torch_geometric.data import Batch
|
|
|
|
x = [torch.rand(100, 6) for _ in range(10)]
|
|
pos = [torch.rand(100, 3) for _ in range(10)]
|
|
graph = [
|
|
KNNGraph(x=x_, pos=pos_, neighbours=6, edge_attr=True)
|
|
for x_, pos_ in zip(x, pos)
|
|
]
|
|
input_ = Batch.from_data_list(graph)
|
|
|
|
|
|
@pytest.mark.parametrize("shared_weights", [True, False])
|
|
def test_constructor(shared_weights):
|
|
lifting_operator = torch.nn.Linear(6, 16)
|
|
projection_operator = torch.nn.Linear(16, 3)
|
|
GraphNeuralOperator(
|
|
lifting_operator=lifting_operator,
|
|
projection_operator=projection_operator,
|
|
edge_features=3,
|
|
internal_layers=[16, 16],
|
|
shared_weights=shared_weights,
|
|
)
|
|
|
|
GraphNeuralOperator(
|
|
lifting_operator=lifting_operator,
|
|
projection_operator=projection_operator,
|
|
edge_features=3,
|
|
inner_size=16,
|
|
internal_n_layers=10,
|
|
shared_weights=shared_weights,
|
|
)
|
|
|
|
int_func = torch.nn.Softplus
|
|
ext_func = torch.nn.ReLU
|
|
|
|
GraphNeuralOperator(
|
|
lifting_operator=lifting_operator,
|
|
projection_operator=projection_operator,
|
|
edge_features=3,
|
|
internal_n_layers=10,
|
|
shared_weights=shared_weights,
|
|
internal_func=int_func,
|
|
external_func=ext_func,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("shared_weights", [True, False])
|
|
def test_forward_1(shared_weights):
|
|
lifting_operator = torch.nn.Linear(6, 16)
|
|
projection_operator = torch.nn.Linear(16, 3)
|
|
model = GraphNeuralOperator(
|
|
lifting_operator=lifting_operator,
|
|
projection_operator=projection_operator,
|
|
edge_features=3,
|
|
internal_layers=[16, 16],
|
|
shared_weights=shared_weights,
|
|
)
|
|
output_ = model(input_)
|
|
assert output_.shape == torch.Size([1000, 3])
|
|
|
|
|
|
@pytest.mark.parametrize("shared_weights", [True, False])
|
|
def test_forward_2(shared_weights):
|
|
lifting_operator = torch.nn.Linear(6, 16)
|
|
projection_operator = torch.nn.Linear(16, 3)
|
|
model = GraphNeuralOperator(
|
|
lifting_operator=lifting_operator,
|
|
projection_operator=projection_operator,
|
|
edge_features=3,
|
|
inner_size=32,
|
|
internal_n_layers=2,
|
|
shared_weights=shared_weights,
|
|
)
|
|
output_ = model(input_)
|
|
assert output_.shape == torch.Size([1000, 3])
|
|
|
|
|
|
@pytest.mark.parametrize("shared_weights", [True, False])
|
|
def test_backward(shared_weights):
|
|
lifting_operator = torch.nn.Linear(6, 16)
|
|
projection_operator = torch.nn.Linear(16, 3)
|
|
model = GraphNeuralOperator(
|
|
lifting_operator=lifting_operator,
|
|
projection_operator=projection_operator,
|
|
edge_features=3,
|
|
internal_layers=[16, 16],
|
|
shared_weights=shared_weights,
|
|
)
|
|
input_.x.requires_grad = True
|
|
output_ = model(input_)
|
|
l = torch.mean(output_)
|
|
l.backward()
|
|
assert input_.x.grad.shape == torch.Size([1000, 6])
|
|
|
|
|
|
@pytest.mark.parametrize("shared_weights", [True, False])
|
|
def test_backward_2(shared_weights):
|
|
lifting_operator = torch.nn.Linear(6, 16)
|
|
projection_operator = torch.nn.Linear(16, 3)
|
|
model = GraphNeuralOperator(
|
|
lifting_operator=lifting_operator,
|
|
projection_operator=projection_operator,
|
|
edge_features=3,
|
|
inner_size=32,
|
|
internal_n_layers=2,
|
|
shared_weights=shared_weights,
|
|
)
|
|
input_.x.requires_grad = True
|
|
output_ = model(input_)
|
|
l = torch.mean(output_)
|
|
l.backward()
|
|
assert input_.x.grad.shape == torch.Size([1000, 6])
|