How to assign a train_mask, test_mask and valid_mask for a PyG graph created from networkx

[RamachandranSitaraman-hv]

Hi,
How to assign a train_mask, test_mask and valid_mask for a PyG graph created from networkx...is there some property of the node i can set and have PyG automatically execute the masks while training, testing and validation? Thanks

[rusty1s]

You can add splitting masks automatically via torch_geometric.transforms.RandomNodeSplit. Let me know if this fixes your issues :)

[RamachandranSitaraman-hv]

In the example you have pointed to, the mask is used only for calculating loss, pred and acc after training, so if i want to limit the training itself only to those nodes defined by the train_mask i would have to implement that myself in the forward function, is that correct understanding? thanks

[rusty1s]

Yes, you can use data.subgraph(data.train_mask) for this.

[RamachandranSitaraman-hv]

Thats great, Could you please point me to an example code illustrating how data.subgraph is used...

[rusty1s]

I don't think we have an example of this. It will simply return a new data object (holding a subset of the data according to train_mask) which you can use just as the original one for GNN training:

train_data = data.subgraph(data.train_mask)
train_out = model(train_data.x, train_data.edge_index)
loss = loss_fn(train_out, train_data.y)

Let me know if something is not clear.

[RamachandranSitaraman-hv]

This is very clear, thank you