[FEAT] Template Fused Gromov Wasserstein layer #488
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
rflamary
changed the title
rflamary
reviewed
Jul 7, 2023
There was a problem hiding this comment.
Thanks @SoniaMaz8 , here are a few comments .
Please also add a quick example of the use of TFGW in the doc.
ot/gnn/_layers.py
Outdated
| n_templates_nodes : int | ||
| Number of nodes in each template. | ||
| alpha0 : float, optional | ||
| Trade-off parameter (0 < alpha < 1). If None alpha is trained, else it is fixed at the given value. |
There was a problem hiding this comment.
FGW tradoff parameter, weights features (apha=) and structures(alpha=).
| """ | ||
| Template Fused Gromov-Wasserstein (TFGW) layer. This layer is a pooling layer for graph neural networks. | ||
| It computes the fused Gromov-Wasserstein distances between the graph and a set of templates. | ||
There was a problem hiding this comment.
add warning that alpha=logitg() tthat is optimized (to avoid constrained optimization)
ot/gnn/_layers.py
Outdated
| self.q0 = nn.Parameter(self.q0) | ||
|
|
||
| if alpha0 is None: | ||
| alpha0 = torch.Tensor([0]) |
There was a problem hiding this comment.
handle alpha as array (add option for that)
ot/gnn/_utils.py
Outdated
| from ..gromov import fused_gromov_wasserstein2 | ||
|
|
||
|
|
||
| def template_initialisation(n_templates, n_template_nodes, n_features, feature_init_mean=0., feature_init_std=1.): |
There was a problem hiding this comment.
TFGW_template_initialiation
test/test_gnn.py
Outdated
| Pooling architecture using the LTFGW layer. | ||
| Parameters | ||
| ---------- |
Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>
Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>
rflamary
changed the title
rflamary
changed the title
rflamary
reviewed
Jul 19, 2023
There was a problem hiding this comment.
Hello @SoniaMaz8,
The PR is looking very good. I have a few comments below and more high level stuff that needs be done before merging.
- Add
gnnin thedocs/source/all.rstfile so that the sub-module is properly documented (it might show some worng formating that needs to be corected) - Express what is computed in math form when describing the Pooling layer, it will help users understand what is done without the need for reading the papers (seen ot.emd2 as an example of how to add math in doc format).
- Add an item in the itemize in the top of the readme such as : Graph Neural Network OT layers TFGW [52] and TW (OT-GNN) [53]
ot/gnn/_layers.py
Outdated
| References | ||
| ---------- | ||
| .. [53] Bécigneul, G., Ganea, O. E., Chen, B., Barzilay, R., & Jaakkola, T. S. (2020). [Optimal transport graph neural networks] (https://arxiv.org/pdf/2006.04804) |
There was a problem hiding this comment.
bad formatting when compiled in doc, you can remove the link.
|
|
||
|
|
||
| @pytest.mark.skipif(not torch_geometric, reason="pytorch_geometric not installed") | ||
| def test_TFGW(): |
There was a problem hiding this comment.
split this in two tests please
Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Types of changes
I introduced a new module ot.gnn in which the Template Fused Gromov Wasserstein layer (Vincent-Cuaz, C., Flamary, R., Corneli, M., Vayer, T., & Courty, N. (2022). Template-based graph neural network with optimal transport distances. Advances in Neural Information Processing Systems, 35, 11800-11814.) is coded. This layer is a pooling layer for GNNs that constructs graph embeddings by constructing a vector of FGW distances between the graph and learned graph templates.
Motivation and context / Related issue
How has this been tested (if it applies)
I created a new test which constructs two small graphs, passes them through the layer and performs back propagation.
PR checklist