| 164 | |
| 165 | @pytest.mark.skipif(not torch_geometric, reason="pytorch_geometric not installed") |
| 166 | def test_TW_variants(): |
| 167 | # Test the TFGW layer by passing two graphs through the layer and doing backpropagation. |
| 168 | |
| 169 | class GNN_pooling(nn.Module): |
| 170 | """ |
| 171 | Pooling architecture using the TW layer. |
| 172 | """ |
| 173 | |
| 174 | def __init__(self, n_features, n_templates, n_template_nodes, pooling_layer): |
| 175 | """ |
| 176 | Pooling architecture using the TW layer. |
| 177 | """ |
| 178 | super().__init__() |
| 179 | |
| 180 | self.n_features = n_features |
| 181 | self.n_templates = n_templates |
| 182 | self.n_template_nodes = n_template_nodes |
| 183 | |
| 184 | self.TFGW = pooling_layer |
| 185 | |
| 186 | self.linear = Linear(self.n_templates, 1) |
| 187 | |
| 188 | def forward(self, x, edge_index, batch=None): |
| 189 | x = self.TFGW(x, edge_index, batch=batch) |
| 190 | |
| 191 | x = self.linear(x) |
| 192 | |
| 193 | return x |
| 194 | |
| 195 | n_templates = 3 |
| 196 | n_template_nodes = 3 |
| 197 | n_nodes = 10 |
| 198 | n_features = 3 |
| 199 | |
| 200 | torch.manual_seed(0) |
| 201 | |
| 202 | C1 = torch.randint(0, 2, size=(n_nodes, n_nodes)) |
| 203 | edge_index1 = torch.stack(torch.where(C1 == 1)) |
| 204 | x1 = torch.rand(n_nodes, n_features) |
| 205 | graph1 = GraphData(x=x1, edge_index=edge_index1, y=torch.tensor([0.0])) |
| 206 | batch1 = torch.tensor([1] * n_nodes) |
| 207 | batch1[: n_nodes // 2] = 0 |
| 208 | |
| 209 | criterion = torch.nn.CrossEntropyLoss() |
| 210 | |
| 211 | for train_node_weights in [True, False]: |
| 212 | model = GNN_pooling( |
| 213 | n_features, |
| 214 | n_templates, |
| 215 | n_template_nodes, |
| 216 | pooling_layer=TWPooling( |
| 217 | n_templates, |
| 218 | n_template_nodes, |
| 219 | n_features, |
| 220 | train_node_weights=train_node_weights, |
| 221 | ), |
| 222 | ) |
| 223 | |