How to specifically train and test with some of nodes in the graph for node classification? Each graph uses different nodes.

[KevinHooah]

Hi, I am working on a node classification task. I have a dataset of 1000 graphs, each graph is an 11*7 grid graph. Each node in each graph has three different labels (0,1,2). Only nodes with labels 0 or 1 are interested. Roughly, each graph has 70 nodes with label 2 (which we are not interested in). I want to train a node classification model with only nodes having labels 0 or 1 in each graph, on 800 graphs, and test on the rest 200 graphs. The node of interest in each graph is not consistent (e.g., graph A may have nodes 1 and 2, while in graph B, they are nodes 2,4,5). I want to ask how to build the dataset with this setting and train a (for example) GCN model.

For now, I am using a method to add the mask as the following examples:

import torch
from torch_geometric.data import Data

# Example graph data
graphs = [Data(x=torch.tensor([[-1], [0], [1]]), edge_index=torch.tensor([[0, 1], [1, 0]])),
          Data(x=torch.tensor([[0], [1], [2]]), edge_index=torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]]))]

train_nodes_list = [torch.tensor([0, 2], dtype=torch.long), torch.tensor([1], dtype=torch.long)]

train_data_list = []
for graph, train_nodes in zip(graphs, train_nodes_list):
    # Create a mask for the selected nodes
    train_mask = torch.zeros(graph.num_nodes, dtype=torch.bool)
    train_mask[train_nodes] = 1

    # Create a subset of the data to only include the selected nodes
    train_data = graph.__class__(x=graph.x[train_mask], edge_index=graph.edge_index)
    train_data_list.append(train_data)

# Use the train_data_list in the training loop
train_node_dataloader = torch_geometric.loader.DataLoader(train_data_list, batch_size=BATCH_SIZE, shuffle=True)

And test with GCN mode like this:

class GCN(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        # torch.manual_seed(1234567)
        self.conv1 = GCNConv(1, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
        self.conv3 = GCNConv(hidden_channels, hidden_channels//2)
        self.conv4 = GCNConv(hidden_channels//2, 2)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = x.relu()
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        x = x.relu()
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv3(x, edge_index)
        x = x.relu()
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv4(x, edge_index)

        return x

However, I am having an index out-of-bounds error which can be traced to \torch_scatter\scatter.py:21, in scatter_sum(src, index, dim, out, dim_size)

Can anyone show me the proper way to build the dataset and train the model? Thank you!

[rusty1s]

The problem is in line

 train_data = graph.__class__(x=graph.x[train_mask], edge_index=graph.edge_index)

where you reduce the number of nodes but keep the original edge_index.

If you want to remove nodes, it's better to use data.subgraph(train_mask) instead.

[KevinHooah]

Thank you very much! Does this method model the problem as a node classification for heterogeneous graphs? I am new to the Graph Learning field. May I ask if there are any suggestions/methods/papers on tackling the task I describe? Thank you very much!

[rusty1s]

The code snippet you showed models a GNN on a homogeneous graph (single node type, single edge type). You can look here on how to get started with heterogeneous graphs.