Concat node features with graph features

[tferber]

My graph consists of nodes with some features each.
I would like to add the mean of each feature per graph (via global_mean_pool(batch.x, batch.batch)) to each node.
I could use to_dense_batch and append them there, but how do i get back the sparse version afterwards?

[rusty1s]

I suggest to use dataset transforms for that, which allows to just add the global node feature to each graph in separation:

def my_transform(data):
    mean = data.x.mean(dim=0, keepdim=True)
    data.x = torch.cat([x, mean], dim=0)
    return data

Dataset(..., transform=my_transform)

[tferber]

Thanks! That would work for the initial inputs, but what if i want to that later, say after the initial feature vectors have been replaced/expanded with some feature transformations?

[rusty1s]

In that case, you have to do it batch-wise. However, node features in the batch do not have be sorted: You can just append them to the list of node features:

out = global_mean_pool(batch.x, batch.batch)
batch.x = torch.cat([batch.x, out, dim=0)
batch.batch = torch.cat([batch.batch, torch.arange(batch.num_graphs)], dim=0)

[tferber]

But this appends them as additional nodes in the batch, no? Maybe i was not clear, sorry.

Here is an example where I have two graphs in the batch, 53 nodes in total, three input features for each node:

torch.Size([53, 3])
tensor([[ 5.3611e-02,  8.7658e-02,  1.7413e-03],
        [ 5.3639e-02,  1.3106e-01,  8.6554e-02],
        [ 5.3537e-02,  1.7496e-01,  5.8730e-03],
        [ 5.3566e-02,  2.1836e-01,  4.7926e-03],
        [ 8.5247e-02,  4.4746e-04,  1.2054e-03],
        [ 8.5033e-02,  4.3843e-02,  5.6456e-03],
        [ 8.5171e-02,  8.7737e-02,  1.1376e-02],
        [ 8.4958e-02,  1.3114e-01,  7.2176e-01],
  ...
        [-1.7718e-02,  1.7500e-01,  2.3416e-03],
        [ 1.7379e-02, -1.7438e-01,  1.2541e-03],
        [ 1.7143e-02, -1.3098e-01,  1.1486e-03],
        [ 1.7096e-02, -4.3643e-02,  1.2056e-03],
        [ 1.7247e-02,  8.7657e-02,  4.1676e-03],
        [ 1.7019e-02,  1.3106e-01,  1.2311e-02],
        [ 1.7215e-02,  1.7501e-01,  3.1643e-03],
        [ 5.3018e-02,  8.7686e-02,  1.7349e-02],
        [ 5.2820e-02,  1.3110e-01,  6.6449e-01],
        [ 5.2986e-02,  1.7504e-01,  2.8501e-03]])

If i run the code by you i get this

torch.Size([55, 3])
tensor([[ 5.3611e-02,  8.7658e-02,  1.7413e-03],
        [ 5.3639e-02,  1.3106e-01,  8.6554e-02],
        [ 5.3537e-02,  1.7496e-01,  5.8730e-03],
        [ 5.3566e-02,  2.1836e-01,  4.7926e-03],
        [ 8.5247e-02,  4.4746e-04,  1.2054e-03],
        [ 8.5033e-02,  4.3843e-02,  5.6456e-03],
        [ 8.5171e-02,  8.7737e-02,  1.1376e-02],
        [ 8.4958e-02,  1.3114e-01,  7.2176e-01],
...
        [-1.7718e-02,  1.7500e-01,  2.3416e-03],
        [ 1.7379e-02, -1.7438e-01,  1.2541e-03],
        [ 1.7143e-02, -1.3098e-01,  1.1486e-03],
        [ 1.7096e-02, -4.3643e-02,  1.2056e-03],
        [ 1.7247e-02,  8.7657e-02,  4.1676e-03],
        [ 1.7019e-02,  1.3106e-01,  1.2311e-02],
        [ 1.7215e-02,  1.7501e-01,  3.1643e-03],
        [ 5.3018e-02,  8.7686e-02,  1.7349e-02],
        [ 5.2820e-02,  1.3110e-01,  6.6449e-01],
        [ 5.2986e-02,  1.7504e-01,  2.8501e-03],
        [ 1.7048e-02,  3.2063e-02,  4.0314e-02],
        [ 7.3091e-03, -1.4691e-03,  3.5401e-02]])

What i want to get is

torch.Size([53, 6])
tensor([[ 5.3611e-02,  8.7658e-02,  1.7413e-03, 1.7048e-02,  3.2063e-02,  4.0314e-02],
        [ 5.3639e-02,  1.3106e-01,  8.6554e-02, 1.7048e-02,  3.2063e-02,  4.0314e-02],
        . . . 
        [ 5.2820e-02,  1.3110e-01,  6.6449e-01,  7.3091e-03, -1.4691e-03,  3.5401e-02],
        [ 5.2986e-02,  1.7504e-01,  2.8501e-03,  7.3091e-03, -1.4691e-03,  3.5401e-02]

Maybe one can think of this as a global message exchange with all nodes in the graph.

[rusty1s]

I see:

out = global_mean_pool(batch.x, batch.batch)
batch.x = torch.cat([batch.x, out[batch]], dim=-1)

should do the trick.

[tferber]

It does the trick, thank you !