Pooling architecture using the TFGW layer.
| 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 |
no outgoing calls