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