Pooling architecture using the TFGW layer.
| 127 | |
| 128 | |
| 129 | class pooling_TFGW(nn.Module): |
| 130 | """ |
| 131 | Pooling architecture using the TFGW layer. |
| 132 | """ |
| 133 | |
| 134 | def __init__( |
| 135 | self, |
| 136 | n_features, |
| 137 | n_templates, |
| 138 | n_template_nodes, |
| 139 | n_classes, |
| 140 | n_hidden_layers, |
| 141 | feature_init_mean=0.0, |
| 142 | feature_init_std=1.0, |
| 143 | ): |
| 144 | """ |
| 145 | Pooling architecture using the TFGW layer. |
| 146 | """ |
| 147 | super().__init__() |
| 148 | |
| 149 | self.n_templates = n_templates |
| 150 | self.n_template_nodes = n_template_nodes |
| 151 | self.n_hidden_layers = n_hidden_layers |
| 152 | self.n_features = n_features |
| 153 | |
| 154 | self.conv = GCNConv(self.n_features, self.n_hidden_layers) |
| 155 | |
| 156 | self.TFGW = TFGWPooling( |
| 157 | self.n_hidden_layers, |
| 158 | self.n_templates, |
| 159 | self.n_template_nodes, |
| 160 | feature_init_mean=feature_init_mean, |
| 161 | feature_init_std=feature_init_std, |
| 162 | ) |
| 163 | |
| 164 | self.linear = Linear(self.n_templates, n_classes) |
| 165 | |
| 166 | def forward(self, x, edge_index, batch=None): |
| 167 | x = self.conv(x, edge_index) |
| 168 | |
| 169 | x = self.TFGW(x, edge_index, batch) |
| 170 | |
| 171 | x_latent = x # save latent embeddings for visualization |
| 172 | |
| 173 | x = self.linear(x) |
| 174 | |
| 175 | return x, x_latent |
| 176 | |
| 177 | |
| 178 | ############################################################################## |