[Roadmap] `torch_geometric.nn.aggr` 🚀
See original GitHub issue🚀 The feature, motivation and pitch
The goal of this roadmap is to unify the concepts of aggregation inside GNNs across both MessagePassing
and global readouts. Currently, these concepts are separated, e.g., via MessagePassing.aggr = "mean"
and global_mean_pool(...)
while the underlying implementation is the same. In addition, some aggregations are only available as global pooling operators (global_sort_pool
, Set2Set
, …), while, in theory, they are also applicable during MessagePassing
(and vice versa, e.g., SAGEConv.aggr = "lstm"
). One additional feature is the combination of aggregations, which is a useful feature both in MessagePassing
(PNAConv
, EGConv
, …) and global readouts.
As such, we want to provide re-usable aggregations as part of a newly defined torch_geometric.nn.aggr.*
package. Unifying these concepts also helps us to perform optimization and specialized implementations in a single place (e.g., fused kernels for multiple aggregations). After integration, the following functionality is applicable:
class MyConv(MessagePassing):
def __init__(self):
super().__init__(aggr="mean")
class MyConv(MessagePassing):
def __init__(self):
super().__init__(aggr=LSTMAggr(channels=...))
class MyConv(MessagePassing):
def __init__(self):
super().__init__(aggr=MultiAggr("mean", "max", Set2Set(channels=...))
Roadmap
The general roadmap looks as follows (at best, each implemented in a separate PR):
- Define
torch_geometric.nn.aggr.*
and implement aBaseAggr
abstract class (#4687) - Add new aggregators:
- Allow for multiple aggregations as part of a
MultiAggr
class (#4749) - Add support for
class-resolver
, similar to here (#4749, #4716) - Ensure
torch.jit.script
support (#4779) - Integrate new aggregations into
MessagePassing
interface (#4779) - Move aggregators from
torch_geometric.nn.glob
totorch_geometric.nn.aggr
(respecting the new interface), deprecate old implementation:-
MeanAggr
,SumAggr
,MaxAggr
,MinAggr
,MulAggr
,VarAggr
,StdAggr
(#4687, #4749) -
MedianAggr
(#5098) -
AttentionalAggr
(#4986) -
Set2Set
(#4762) -
GlobalSortAggr
(#4957) -
GraphMultiSetTransformer
(#4973) -
EquilibriumAggr
(#4522) - Deprecate
torch_geometric.nn.glob
(#5039)
-
- Update existing GNN layers to make use of new interface, e.g.:
-
SAGEConv
(#4863) -
PNAConv
(#4864) -
GravNetConv
(#4865) -
GENConv
(#4866)
-
- Add support for “reverse” aggregation resolver to keep
message_and_aggregate
functionality intact (#5084) - Support for multiple aggregations in
SAGEConv
(#5033) -
MultiAggregation
: Support forconcat
,concat+transform
,sum
,mean
,max
,attention
(#5000, #5034) - Add
semi_grad
functionality toSoftmaxAggregation
(#4995) - Add and verify
torch_geometric.nn.aggr.*
documentation (#5036, #5097, #5099, #5104) - Add a tutorial on the new concepts (#4927)
- Kernel fusion: Optimize aggregations, e.g., by computing multiple aggregations in parallel (at best discussed in a separate issue)
Any feedback and help from the community is highly appreciated!
cc: @lightaime @Padarn
Issue Analytics
- State:
- Created a year ago
- Reactions:6
- Comments:6 (6 by maintainers)
Thanks everyone for the hard work @lightaime @Padarn. I think that the final outcome looks fantastic - many cool things to promote in our upcoming release 😃
Yes, indeed. There do not exist clear plans for implementation yet though. It will likely depend on PyTorch fusing this ops as part of TorchScript, or on us providing special CUDA kernels.