Graph Pooling Framework :rocket:

rusty1s commented 1 year ago

🚀 The feature, motivation and pitch

The "Understanding Pooling in Graph Neural Networks" paper introduces a simple framework to unify (most) graph pooling approaches:

Select: Selects input nodes to map to supernodes
Reduce: Reduces the supernodes to singletons
Connect: Decides how the new nodes are connected

Currently, the pooling operators included in PyG are implemented isolated from each other, and there exists a lot of repetitive code and different interfaces, which makes their application confusing and challenging. By following the above approach, we can unify existing implementations and can accelerate new research on graph pooling. For example, we can introduce base classes for each of the aforementioned steps:

class Select(torch.nn.Module):
    def forward(self, *args, **kwargs) -> torch.Tensor:
         """Returns a bipartite `edge_index` mapping input nodes to supernodes."""

class Reduce:  # No need -> we can utilize `nn.aggr` for this
    pass

class Connect(torch.nn.Module):
    def forward(self, cluster_index: torch.Tensor, edge_index: torch.Tensor, *args, **kwargs):
        """Returns a coarsened graph."""

With this, e.g., graclus can be moved to a Select operator and TopK pooling can be moved to a Connect operator.

Relevant twitter thread: https://twitter.com/riceasphait/status/1447867635442585601

Tasks:

[x] Add Select and Connect base classes [#6540]
[x] Add Pooling base class [#6540]

Add concrete Select and Connect classes and update implementations of.

[x] TopKPooling [#7308]
[x] #7617

wsad1 commented 1 year ago

Would it be useful to add a Pooling class that all pooling operators inherit from.

class Pooling(torch.nn.Module):
    def forward(self, x, edge_index, *args, **kwargs):
          mapping = self.select(x, edge_index, *args, **kwargs)
          x = self.reduce(x, mapping)
          edge_index = self.connect(mapping, edge_index, *args, **kwargs)

self.select, reduce and connect would be objects defined in the child class like graclus or TopK.

rusty1s commented 1 year ago

Yes, definitely :)

danielegrattarola commented 1 year ago

Hey, first author of the paper here!

This is super, if help is still wanted (as I see on the tags) I'd be happy to participate!

Worth mentioning that there is an implementation of SRC in Spektral and a few layers implemented with it, it might be possible to translate the implementation from TF to Torch/PyG.

wsad1 commented 1 year ago

@danielegrattarola firstly great paper. Sorry for the late reply. But if you are still interested we could use your help moving SAGPooling and PANPooling. Just refactor topkpooling here.