Pooling architecture using the TW layer.
| 234 | # Test the TW layer by passing two graphs through the layer and doing backpropagation. |
| 235 | |
| 236 | class pooling_TW(nn.Module): |
| 237 | """ |
| 238 | Pooling architecture using the TW layer. |
| 239 | """ |
| 240 | |
| 241 | def __init__(self, n_features, n_templates, n_template_nodes): |
| 242 | """ |
| 243 | Pooling architecture using the TW layer. |
| 244 | """ |
| 245 | super().__init__() |
| 246 | |
| 247 | self.n_features = n_features |
| 248 | self.n_templates = n_templates |
| 249 | self.n_template_nodes = n_template_nodes |
| 250 | |
| 251 | self.TFGW = TWPooling( |
| 252 | self.n_templates, self.n_template_nodes, self.n_features |
| 253 | ) |
| 254 | |
| 255 | self.linear = Linear(self.n_templates, 1) |
| 256 | |
| 257 | def forward(self, x, edge_index): |
| 258 | x = self.TFGW(x, edge_index) |
| 259 | |
| 260 | x = self.linear(x) |
| 261 | |
| 262 | return x |
| 263 | |
| 264 | torch.manual_seed(0) |
| 265 |