Time complexity for pooling layers

[davidireland3]

Has anyone experimented with any of the differentiable pooling layers to try and optimise for a speed?

I have been using TopKPooling and noticed a considerable speed up when removing it, but my GNN's are unable to generalise as well when not using it. I was wondering if anyone has experimented with any of the other pooling/coarsening layers available in PyG to compare with speed?
When I originally chose to use the k-pooling it is because it seemed quite a lightweight operation, so I would be interested to know if anything faster is available (or if I am totally wrong about it being lightweight).

[rusty1s]

In general, pooling layers involve a lot of computation (computing clusters, filtering node features, coarsening the graph), so it's natural that they can be a bottleneck. Do you have some more information which part exactly becomes the bottleneck?