()
| 22 | |
| 23 | @pytest.mark.skipif(not torch_geometric, reason="pytorch_geometric not installed") |
| 24 | def test_TFGW_optim(): |
| 25 | # Test the TFGW layer by passing two graphs through the layer and doing backpropagation. |
| 26 | |
| 27 | class pooling_TFGW(nn.Module): |
| 28 | """ |
| 29 | Pooling architecture using the TFGW layer. |
| 30 | """ |
| 31 | |
| 32 | def __init__(self, n_features, n_templates, n_template_nodes): |
| 33 | """ |
| 34 | Pooling architecture using the TFGW layer. |
| 35 | """ |
| 36 | super().__init__() |
| 37 | |
| 38 | self.n_features = n_features |
| 39 | self.n_templates = n_templates |
| 40 | self.n_template_nodes = n_template_nodes |
| 41 | |
| 42 | self.TFGW = TFGWPooling( |
| 43 | self.n_templates, self.n_template_nodes, self.n_features |
| 44 | ) |
| 45 | |
| 46 | self.linear = Linear(self.n_templates, 1) |
| 47 | |
| 48 | def forward(self, x, edge_index): |
| 49 | x = self.TFGW(x, edge_index) |
| 50 | |
| 51 | x = self.linear(x) |
| 52 | |
| 53 | return x |
| 54 | |
| 55 | torch.manual_seed(0) |
| 56 | |
| 57 | n_templates = 3 |
| 58 | n_template_nodes = 3 |
| 59 | n_nodes = 10 |
| 60 | n_features = 3 |
| 61 | n_epochs = 3 |
| 62 | |
| 63 | C1 = torch.randint(0, 2, size=(n_nodes, n_nodes)) |
| 64 | C2 = torch.randint(0, 2, size=(n_nodes, n_nodes)) |
| 65 | |
| 66 | edge_index1 = torch.stack(torch.where(C1 == 1)) |
| 67 | edge_index2 = torch.stack(torch.where(C2 == 1)) |
| 68 | |
| 69 | x1 = torch.rand(n_nodes, n_features) |
| 70 | x2 = torch.rand(n_nodes, n_features) |
| 71 | |
| 72 | graph1 = GraphData(x=x1, edge_index=edge_index1, y=torch.tensor([0.0])) |
| 73 | graph2 = GraphData(x=x2, edge_index=edge_index2, y=torch.tensor([1.0])) |
| 74 | |
| 75 | dataset = DataLoader([graph1, graph2], batch_size=1) |
| 76 | |
| 77 | model_FGW = pooling_TFGW(n_features, n_templates, n_template_nodes) |
| 78 | |
| 79 | optimizer = torch.optim.Adam(model_FGW.parameters(), lr=0.01) |
| 80 | criterion = torch.nn.CrossEntropyLoss() |
| 81 |
nothing calls this directly
no test coverage detected