Function test_TW

test/test_gnn.py:233–299

()

@pytest.mark.skipif(not torch_geometric, reason="pytorch_geometric not installed")
def test_TW():
    # Test the TW layer by passing two graphs through the layer and doing backpropagation.

    class pooling_TW(nn.Module):
        """
        Pooling architecture using the TW layer.
        """

        def __init__(self, n_features, n_templates, n_template_nodes):
            """
            Pooling architecture using the TW layer.
            """
            super().__init__()

            self.n_features = n_features
            self.n_templates = n_templates
            self.n_template_nodes = n_template_nodes

            self.TFGW = TWPooling(
                self.n_templates, self.n_template_nodes, self.n_features
            )

            self.linear = Linear(self.n_templates, 1)

        def forward(self, x, edge_index):
            x = self.TFGW(x, edge_index)

            x = self.linear(x)

            return x

    torch.manual_seed(0)

    n_templates = 3
    n_template_nodes = 3
    n_nodes = 10
    n_features = 3
    n_epochs = 3

    C1 = torch.randint(0, 2, size=(n_nodes, n_nodes))
    C2 = torch.randint(0, 2, size=(n_nodes, n_nodes))

    edge_index1 = torch.stack(torch.where(C1 == 1))
    edge_index2 = torch.stack(torch.where(C2 == 1))

    x1 = torch.rand(n_nodes, n_features)
    x2 = torch.rand(n_nodes, n_features)

    graph1 = GraphData(x=x1, edge_index=edge_index1, y=torch.tensor([0.0]))
    graph2 = GraphData(x=x2, edge_index=edge_index2, y=torch.tensor([1.0]))

    dataset = DataLoader([graph1, graph2], batch_size=1)

    model_W = pooling_TW(n_features, n_templates, n_template_nodes)

    optimizer = torch.optim.Adam(model_W.parameters(), lr=0.01)
    criterion = torch.nn.CrossEntropyLoss()