()
| 92 | |
| 93 | @pytest.mark.skipif(not torch_geometric, reason="pytorch_geometric not installed") |
| 94 | def test_TFGW_variants(): |
| 95 | # Test the TFGW layer by passing two graphs through the layer and doing backpropagation. |
| 96 | |
| 97 | class GNN_pooling(nn.Module): |
| 98 | """ |
| 99 | Pooling architecture using the TW layer. |
| 100 | """ |
| 101 | |
| 102 | def __init__(self, n_features, n_templates, n_template_nodes, pooling_layer): |
| 103 | """ |
| 104 | Pooling architecture using the TW layer. |
| 105 | """ |
| 106 | super().__init__() |
| 107 | |
| 108 | self.n_features = n_features |
| 109 | self.n_templates = n_templates |
| 110 | self.n_template_nodes = n_template_nodes |
| 111 | |
| 112 | self.TFGW = pooling_layer |
| 113 | |
| 114 | self.linear = Linear(self.n_templates, 1) |
| 115 | |
| 116 | def forward(self, x, edge_index, batch=None): |
| 117 | x = self.TFGW(x, edge_index, batch=batch) |
| 118 | |
| 119 | x = self.linear(x) |
| 120 | |
| 121 | return x |
| 122 | |
| 123 | n_templates = 3 |
| 124 | n_template_nodes = 3 |
| 125 | n_nodes = 10 |
| 126 | n_features = 3 |
| 127 | |
| 128 | torch.manual_seed(0) |
| 129 | |
| 130 | C1 = torch.randint(0, 2, size=(n_nodes, n_nodes)) |
| 131 | edge_index1 = torch.stack(torch.where(C1 == 1)) |
| 132 | x1 = torch.rand(n_nodes, n_features) |
| 133 | graph1 = GraphData(x=x1, edge_index=edge_index1, y=torch.tensor([0.0])) |
| 134 | batch1 = torch.tensor([1] * n_nodes) |
| 135 | batch1[: n_nodes // 2] = 0 |
| 136 | |
| 137 | criterion = torch.nn.CrossEntropyLoss() |
| 138 | |
| 139 | for train_node_weights in [True, False]: |
| 140 | for alpha in [None, 0, 0.5]: |
| 141 | for multi_alpha in [True, False]: |
| 142 | model = GNN_pooling( |
| 143 | n_features, |
| 144 | n_templates, |
| 145 | n_template_nodes, |
| 146 | pooling_layer=TFGWPooling( |
| 147 | n_templates, |
| 148 | n_template_nodes, |
| 149 | n_features, |
| 150 | alpha=alpha, |
| 151 | multi_alpha=multi_alpha, |
nothing calls this directly
no test coverage detected