Node-level regression problem: take inputs of a subset of nodes to predict the behavior of all nodes

[DrKleiber]

I'm new to graph models and am working on a problem for energy system prediction. I model the system as a graph with each node representing a certain component. I can place sensors in certain nodes but not all. So in real life, I can only know the behavior of a subset of nodes of the graph through sensor measurements.

My goal is to train a (static) graph model that reads in the sensor measurements of a subset of nodes and predicts the behavior of all nodes. To do that, I ran a few hundred simulations and get the behavior of all nodes under certain conditions. I want to perform the training based on these results, so the model can be used for unseen conditions. But I didn't find a similar example to my problem.

I read the Colab tutorial for node classification, where we define train_mask so that model reads in the features of all the nodes and calculates the error function based on the labels of a subset of nodes. Looks like contrary to my case, but when I tried to apply the train_mask to input( data.x[data.train_mask]), error showed up, as looks like x[data.train_mask] is inconsistent with the given graph. I wonder is there any possible solution to my problem? Would padding the unknown node with zero in the input be a viable option?

Any suggestions/examples would be highly appreciated! Thank you in advance!

[rusty1s]

Shapes of x and train_mask need to match in the first dimension, that is, the both should hold values for every node in the graph. If you only want to learn on a subgraph but apply the model on the full graph later on, how about you generate your sub-data via

train_data = data.subgraph(train_mask_or_indices)

and use it to train your model in a fully inductive mode. WDYT?

[rusty1s]

This looks like a classical masked autoencoder approach. That is, you input your graph and features into the GNN, where the features for unknown nodes are set to zero. After message passing completes, you try to reconstruct the features of the unknown nodes (using MSE loss):

train_mask = ... # zero for nodes we know the features of.

x = data.x.clone()
x[~train_mask] = 0.
out = model(x, data.edge_index)  # Maps to the same shape as `data.x`
loss = F.mse_loss(out[~train_mask], data.x[~train_mask])
loss.backward()

[DrKleiber]

Thank you @rusty1s ! I'll try it.

[atello]

Hi @rusty1s,

First of all, thanks for helping us, who are starting with GNNs, with your guidance.
I am working on a similar setting than the one that is described here. Just to be sure, in the line where you calculate the mse_loss shouldn't it be out[~train_mask] instead of x[~train_mask]

loss = F.mse_loss(out[~train_mask], data.x[~train_mask])

[rusty1s]

Yes, absolutely. Thanks for pointing that out. Edited my code snippet.

[atello]

Thanks for your reply @rusty1s!