Reproducible global pooling operation?

[realCrush]

Hi, I use PyG to implement my graph pooling model, however, it seems PyG suffers from numerical unstable on GPUs when using torch_scatter based methods, I followed #2788 (comment) to fix the MessagePassing operation with adjacent matrix in SparseTensor format, it works pretty well.

However when I need to implement a graph pooling layer, I need some global pooling operations like global_mean_pool, global_max_pool, global_add_pool , I can't find reproducible alternatives in PyG, I want to ask if PyG provides such alternatives? If not, how can I make them(global_mean_pool, global_max_pool, global_add_pool) reproducible?

[rusty1s]

You can use segment_csr from torch-scatter:

from torch_scatter import segment_csr

global_x = segment_csr(x, batch.ptr, reduce='sum')

[realCrush]

Thank you, I'll check this.

[realCrush]

BTW, is there a time-efficient way to generate ptr tensor from ordered batch tansor? without for-loop.

[rusty1s]

The mini-batch will auto-generate it for you. Alternatively, you can do:

deg = torch_geometric.utils.degree(batch)
ptr = deg.cumsum(dim=0)
ptr = torch.cat([torch.tensor([0]), ptr], dim=0)

[realCrush]

Perfect, thanks!!