r""" Template Wasserstein (TW) layer, also known as OT-GNN layer. This layer is a pooling layer for graph neural networks. Computes the Wasserstein distances between the features of the graph features and a set of templates. .. math:: TW_{\overline{\mathcal{G
(
self,
n_features,
n_tplt=2,
n_tplt_nodes=2,
train_node_weights=True,
feature_init_mean=0.0,
feature_init_std=1.0,
)
| 215 | """ |
| 216 | |
| 217 | def __init__( |
| 218 | self, |
| 219 | n_features, |
| 220 | n_tplt=2, |
| 221 | n_tplt_nodes=2, |
| 222 | train_node_weights=True, |
| 223 | feature_init_mean=0.0, |
| 224 | feature_init_std=1.0, |
| 225 | ): |
| 226 | r""" |
| 227 | Template Wasserstein (TW) layer, also known as OT-GNN layer. This layer is a pooling layer for graph neural networks. |
| 228 | Computes the Wasserstein distances between the features of the graph features and a set of templates. |
| 229 | |
| 230 | .. math:: |
| 231 | TW_{\overline{\mathcal{G}}}(C,F,h)=[W(F,h,\overline{F}_k,\overline{h}_k)]_{k=1}^{K} |
| 232 | |
| 233 | where : |
| 234 | |
| 235 | - :math:`\mathcal{G}=\{(\overline{F}_k,\overline{h}_k) \}_{k \in \llbracket 1;K \rrbracket }` is the set of :math:`K` templates characterized by their feature matrices :math:`\overline{F}_k` and their node weights :math:`\overline{h}_k`. |
| 236 | - :math:`F` and :math:`h` are respectively the feature matrix and the node weights of the graph. |
| 237 | |
| 238 | Parameters |
| 239 | ---------- |
| 240 | n_features : int |
| 241 | Feature dimension of the nodes. |
| 242 | n_tplt : int |
| 243 | Number of graph templates. |
| 244 | n_tplt_nodes : int |
| 245 | Number of nodes in each template. |
| 246 | train_node_weights : bool, optional |
| 247 | If True, the templates node weights are learned. |
| 248 | Else, they are uniform. |
| 249 | feature_init_mean: float, optional |
| 250 | Mean of the random normal law to initialize the template features. |
| 251 | feature_init_std: float, optional |
| 252 | Standard deviation of the random normal law to initialize the template features. |
| 253 | |
| 254 | References |
| 255 | ---------- |
| 256 | .. [54] Bécigneul, G., Ganea, O. E., Chen, B., Barzilay, R., & Jaakkola, T. S. (2020). [Optimal transport graph neural networks] |
| 257 | """ |
| 258 | super().__init__() |
| 259 | |
| 260 | self.n_tplt = n_tplt |
| 261 | self.n_tplt_nodes = n_tplt_nodes |
| 262 | self.n_features = n_features |
| 263 | self.feature_init_mean = feature_init_mean |
| 264 | self.feature_init_std = feature_init_std |
| 265 | |
| 266 | _, tplt_features, self.q0 = TFGW_template_initialization( |
| 267 | self.n_tplt, |
| 268 | self.n_tplt_nodes, |
| 269 | self.n_features, |
| 270 | self.feature_init_mean, |
| 271 | self.feature_init_std, |
| 272 | ) |
| 273 | |
| 274 | self.tplt_features = nn.Parameter(tplt_features) |
no test coverage detected