Using an attention pooling for node features

[fedeotto]

I would like to use a softmax-based attention pooling to final node features of graphs. In particular I want to compute attention scores $a_i = \frac{\exp(x_i)}{\sum_{j} \exp(x_j)}$. Then the graph pooling would be obtained by computing $G = \frac{1}{N} \sum_{i}a_i x_i$.
The problem is that I would like to compute softmax scores by ranging j on the nodes of each graph, because attention coefficients should depend only on the other nodes in the same graph and not other nodes that are included in features tensor. The confusion arises from having a features tensor x storing nodes of multiple graphs, given PyG batching functionalitites. How could I do that, using features tensor x and batch tensor?

[ldv1]

Is this

import torch
from torch_geometric.utils import softmax
from torch_geometric.nn.pool import global_mean_pool
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader

class Model(torch.nn.Module):
    
    def __init__(self):
        super().__init__()
        # other vars
        
    def forward(self,x, batch):
        # something useful here
        att = softmax(x, index=batch)
        y = global_mean_pool(x*att, batch=batch)
        return y

dataset = [ Data(x=torch.Tensor([ [1,1],[1, 1.1],[1,0.9] ])), \
            Data(x=torch.Tensor([ [1,1],[1, 1.1] ])) ]
loader = DataLoader(dataset, batch_size=2)
model = Model()
for step, data in enumerate(loader):
    model(data.x, batch=data.batch)

helpful ?

[fedeotto]

Many thanks! That's exactly what I was looking for. I didn't know about torch_geometric.utils.softmax.