Handling Linear layer dimesions with graph sizes across batches

[diamondspark]

Hi everyone,
Seeking guidance with following. The gist of the problem is below
Trying to calculate scalar attention score per cluster of a graph.
I have several large graphs. I cluster each graph into clusters (say 40 cluster per graph). Each cluster has different number of nodes Lets say 64 is largest sized cluster for an arbitrary graph. Each node is defined by 150 dimensional feature vector. Hence, my data becomes (40*n_batches, 64, 150). Using Linear(150,1), I get output scores of shape (40*n_batches, 64). However, I am seeking a scalar score per cluster i.e. (40*n_batches,1).

Now, I am unsure how to use another linear layer on top of this (40*n_batches, 64) to go to (40*n_batches,1) since 64 is not a fixed quantity and varies across batches (graphs). How do I define this Linear layer in my init

[rusty1s]

You need to make use of permutation-invariant pooling operators to go from node-level features/scalars to graph-level ones, e.g., via torch_geometric.nn.global_mean_pool:

x = self.lin1(x)
x = global_mean_pool(x, cluster_index)
x = self.lin2(x)

[diamondspark]

Oh absolutely. That makes sense. Thank you!