Performing global_mean_pool on a batch of data

[fedeotto]

Hello, I'm trying to understand how global_mean_pool layer works, in particular, from the documentation it's not too clear to me what exactly I should assign as batch argument.

If I have this simple Network

class Net(torch.nn.Module):
    def __init__(self):
        
        super(Net, self).__init__()
        self.conv1 = GCNConv(112, 10)
        self.conv2 = GCNConv(10,8)
        
    def forward(self):
        x = self.conv1(data.x, data.edge_index)
        x = self.conv2(x, edge_index)
        return x

and a batch of data

data = Batch(batch=[274], edge_index=[2, 818], ptr=[101], x=[274, 112], y=[100])

How should I perform a global mean pooling on this?

[wsad1]

You can read more about graph batching here. batch is a long tensor which indicates which graph each node belongs to. For example,

batch = torch.tensor([0,0,0,1,1,1])

The above batch vector indicates that the first 3 nodes belong to 0th graph and last 3 to 1st graph.

On your data the global mean pool call would look something like this

    def forward(self):
        x = self.conv1(data.x, data.edge_index)
        x = self.conv2(x, edge_index)
        pool_x = global_mean_pool(x, data.batch)
        return x