Add test supervised solver for graph data and remove PinaBatch

Author: FilippoOlivo

pina/data/dataset.py

@@ -109,28 +109,6 @@ def input_points(self):
         return {k: v["input_points"] for k, v in self.conditions_dict.items()}
 
 
-class PinaBatch(Batch):
-    """
-    Add extract function to torch_geometric Batch object
-    """
-
-    def __init__(self):
-
-        super().__init__(self)
-
-    def extract(self, labels):
-        """
-        Perform extraction of labels on node features (x)
-
-        :param labels: Labels to extract
-        :type labels: list[str] | tuple[str] | str
-        :return: Batch object with extraction performed on x
-        :rtype: PinaBatch
-        """
-        self.x = self.x.extract(labels)
-        return self
-
-
 class PinaGraphDataset(PinaDataset):
 
     def __init__(
@@ -182,12 +160,10 @@ def _divide_batch(self, batch):
         """
         to_return_dict = {}
         to_return_dict["input_points"] = batch
-        if hasattr(batch, "y"):
-            to_return_dict["output_points"] = batch.y
         return to_return_dict
 
     def _base_create_graph_batch_from_list(self, data):
-        batch = PinaBatch.from_data_list(data)
+        batch = Batch.from_data_list(data)
         return batch
 
     def _getitem_dummy(self, idx):

pina/graph.py

@@ -162,6 +162,18 @@ def _preprocess_edge_index(edge_index, undirected):
             edge_index = to_undirected(edge_index)
         return edge_index
 
+    def extract(self, labels):
+        """
+        Perform extraction of labels on node features (x)
+
+        :param labels: Labels to extract
+        :type labels: list[str] | tuple[str] | str
+        :return: Batch object with extraction performed on x
+        :rtype: PinaBatch
+        """
+        self.x = self.x.extract(labels)
+        return self
+
 
 class GraphBuilder:
     """

tests/test_solver/test_supervised_solver.py

@@ -6,7 +6,9 @@
 from pina.solver import SupervisedSolver
 from pina.model import FeedForward
 from pina.trainer import Trainer
+from pina.graph import KNNGraph
 from torch._dynamo.eval_frame import OptimizedModule
+from torch_geometric.nn import GCNConv
 
 
 class LabelTensorProblem(AbstractProblem):
@@ -30,9 +32,72 @@ class TensorProblem(AbstractProblem):
     }
 
 
+x = torch.rand((100, 20, 5))
+pos = torch.rand((100, 20, 2))
+output_ = torch.rand((100, 20, 1))
+input_ = [
+    KNNGraph(x=x_, pos=pos_, neighbours=3, edge_attr=True, y=y_)
+    for x_, pos_, y_ in zip(x, pos, output_)
+]
+
+
+class GraphProblem(AbstractProblem):
+    output_variables = None
+    conditions = {
+        "data": Condition(
+            graph=input_,
+        )
+    }
+
+
+x = LabelTensor(torch.rand((100, 20, 5)), ["a", "b", "c", "d", "e"])
+pos = LabelTensor(torch.rand((100, 20, 2)), ["x", "y"])
+output_ = LabelTensor(torch.rand((100, 20, 1)), ["u"])
+input_ = [
+    KNNGraph(x=x[i], pos=pos[i], neighbours=3, edge_attr=True, y=output_[i])
+    for i in range(len(x))
+]
+
+
+class GraphProblemLT(AbstractProblem):
+    output_variables = ["u"]
+    input_variables = ["a", "b", "c", "d", "e"]
+    conditions = {
+        "data": Condition(
+            graph=input_,
+        )
+    }
+
+
 model = FeedForward(2, 1)
 
 
+class Model(torch.nn.Module):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.lift = torch.nn.Linear(5, 10)
+        self.activation = torch.nn.Tanh()
+        self.output = torch.nn.Linear(10, 1)
+
+        self.conv = GCNConv(10, 10)
+
+    def forward(self, batch):
+
+        x = batch.x
+        edge_index = batch.edge_index
+        for _ in range(1):
+            y = self.lift(x)
+            y = self.activation(y)
+            y = self.conv(y, edge_index)
+            y = self.activation(y)
+            y = self.output(y)
+        return y
+
+
+graph_model = Model()
+
+
 def test_constructor():
     SupervisedSolver(problem=TensorProblem(), model=model)
     SupervisedSolver(problem=LabelTensorProblem(), model=model)
@@ -64,6 +129,24 @@ def test_solver_train(use_lt, batch_size, compile):
         assert isinstance(solver.model, OptimizedModule)
 
 
+@pytest.mark.parametrize("batch_size", [None, 1, 5, 20])
+@pytest.mark.parametrize("use_lt", [True, False])
+def test_solver_train_graph(batch_size, use_lt):
+    problem = GraphProblemLT() if use_lt else GraphProblem()
+    solver = SupervisedSolver(problem=problem, model=graph_model, use_lt=use_lt)
+    trainer = Trainer(
+        solver=solver,
+        max_epochs=2,
+        accelerator="cpu",
+        batch_size=batch_size,
+        train_size=1.0,
+        test_size=0.0,
+        val_size=0.0,
+    )
+
+    trainer.train()
+
+
 @pytest.mark.parametrize("use_lt", [True, False])
 @pytest.mark.parametrize("compile", [True, False])
 def test_solver_validation(use_lt, compile):
@@ -84,6 +167,24 @@ def test_solver_validation(use_lt, compile):
         assert isinstance(solver.model, OptimizedModule)
 
 
+@pytest.mark.parametrize("batch_size", [None, 1, 5, 20])
+@pytest.mark.parametrize("use_lt", [True, False])
+def test_solver_val_graph(batch_size, use_lt):
+    problem = GraphProblemLT() if use_lt else GraphProblem()
+    solver = SupervisedSolver(problem=problem, model=graph_model, use_lt=use_lt)
+    trainer = Trainer(
+        solver=solver,
+        max_epochs=2,
+        accelerator="cpu",
+        batch_size=batch_size,
+        train_size=0.9,
+        val_size=0.1,
+        test_size=0.0,
+    )
+
+    trainer.train()
+
+
 @pytest.mark.parametrize("use_lt", [True, False])
 @pytest.mark.parametrize("compile", [True, False])
 def test_solver_test(use_lt, compile):
@@ -104,6 +205,24 @@ def test_solver_test(use_lt, compile):
         assert isinstance(solver.model, OptimizedModule)
 
 
+@pytest.mark.parametrize("batch_size", [None, 1, 5, 20])
+@pytest.mark.parametrize("use_lt", [True, False])
+def test_solver_test_graph(batch_size, use_lt):
+    problem = GraphProblemLT() if use_lt else GraphProblem()
+    solver = SupervisedSolver(problem=problem, model=graph_model, use_lt=use_lt)
+    trainer = Trainer(
+        solver=solver,
+        max_epochs=2,
+        accelerator="cpu",
+        batch_size=batch_size,
+        train_size=0.8,
+        val_size=0.1,
+        test_size=0.1,
+    )
+
+    trainer.test()
+
+
 def test_train_load_restore():
     dir = "tests/test_solver/tmp/"
     problem = LabelTensorProblem()
@@ -145,3 +264,52 @@ def test_train_load_restore():
     import shutil
 
     shutil.rmtree("tests/test_solver/tmp")
+
+
+def test_train_load_restore_graph():
+    dir = "tests/test_solver/tmp/"
+    problem = GraphProblemLT()
+    solver = SupervisedSolver(problem=problem, model=graph_model)
+    trainer = Trainer(
+        solver=solver,
+        max_epochs=5,
+        accelerator="cpu",
+        batch_size=None,
+        train_size=0.9,
+        test_size=0.1,
+        val_size=0.0,
+        default_root_dir=dir,
+    )
+    trainer.train()
+
+    # restore
+    new_trainer = Trainer(solver=solver, max_epochs=5, accelerator="cpu")
+    new_trainer.train(
+        ckpt_path=f"{dir}/lightning_logs/version_0/checkpoints/"
+        + "epoch=4-step=5.ckpt"
+    )
+
+    # loading
+    new_solver = SupervisedSolver.load_from_checkpoint(
+        f"{dir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt",
+        problem=problem,
+        model=graph_model,
+    )
+
+    test_pts = KNNGraph(
+        x=LabelTensor(torch.rand(20, 5), ["a", "b", "c", "d", "e"]),
+        pos=LabelTensor(torch.rand(20, 2), ["x", "y"]),
+        neighbours=3,
+        edge_attr=True,
+    )
+
+    assert new_solver.forward(test_pts).shape == (20, 1)
+    assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
+    torch.testing.assert_close(
+        new_solver.forward(test_pts), solver.forward(test_pts)
+    )
+
+    # rm directories
+    import shutil
+
+    shutil.rmtree("tests/test_solver/tmp")