The difference between global_mean_pool in pyg and first reshape then mean in torch

[yeshenpy]

I am a newcomer in the field of GNN, and I have encountered some issues in my implementation.

When processing a batch of data, I need to obtain node representations for each graph and, additionally, obtain graph representations. Currently, I get the representation of each graph using the approach
node_embedding = node_embedding.reshape([batch_size, n_node, feature_dim])
followed by
graph_embedding = node_embedding.mean(1)
I wonder if there is any difference between this approach and using
graph_embedding = global_mean_pool(x=node_embedding , batch=batch)

[rusty1s]

You would want to use global_mean_pool in case your graphs are of different size, in which case you can not simple reshape your node embeddings. If your graphs are of equal size, then the two implementations are the same (while reshaping + torch.mean() should be a bit faster).

[yeshenpy]

I understand, thank you for your response!