Improving node regression performance

[pmoriano]

Hello,
I am working towards a node regression model using GraphSAGE. The graph of interest is unweighted and directed. Node labels are integers. I am using node2vec embeddings as node features. For debugging purposes, I am using a toy graph of 36 nodes and 96 edges. You can find the graph file here.

The below code runs, but the results are not good. I wonder if you can please take a look and provide advise of how to improve the performance of the model. I am not sure if this is because of the small size of the toy graph or if there are issues with my code.

Many thanks.

import numpy as np
import pandas as pd
import networkx as nx

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv
from torch_geometric.nn import SAGEConv

from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from ge import Node2Vec

# Load graph
H = nx.read_graphml("./seed_graph.graphml")

# Extract labels
labels = np.array([H.nodes[str(index)]["label"] for index in range(0, H.number_of_nodes())])

# Compute node features
embedding_model = Node2Vec(H, walk_length=10, num_walks=80, p=0.25, q=4, workers=1)
embedding_model.train(window_size=5, iter=3)
embeddings = embedding_model.get_embeddings()

# Extract edges
H = nx.relabel_nodes(H, {node: int(node) for node in H.nodes})
edges = list(H.edges)

# DataClass 
data = Data(x=torch.from_numpy(np.array(list(embeddings.values()))).type(torch.float), 
            edge_index=torch.tensor(edges).t().contiguous(), y=torch.from_numpy(labels).type(torch.float))

# Splitting dataset
X_train, X_test, y_train, y_test = train_test_split(pd.Series(H.nodes()), pd.Series(labels), test_size=0.30, random_state=7)

n_nodes = H.number_of_nodes()
train_mask = torch.zeros(n_nodes, dtype=torch.bool)
test_mask = torch.zeros(n_nodes, dtype=torch.bool)

train_mask[X_train.index] = True
test_mask[X_test.index] = True

data["train_mask"] = train_mask
data["test_mask"] = test_mask

# Define/Instantiate GNN model
class GNN(torch.nn.Module):
    def __init__(self, in_size, embedding_size, out_size):

        super(GNN, self).__init__()
        self.conv1 = SAGEConv(in_size, embedding_size) # GCNConv, SAGEConv
        self.conv1.aggr = "max"
        self.conv2 = SAGEConv(embedding_size, out_size)
        self.conv2.aggr = "max"

    def forward(self, x, edge_index):

        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        x = F.relu(x)
        
        return x

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model, data = GNN(in_size=data.x.shape[1], embedding_size=16, out_size=1).to(device), data.to(device) 

print(model)
print("Model parameters: ", sum(p.numel() for p in model.parameters()))

# Train/Test Definition
torch.manual_seed(7)

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

def train():
    model.train()
    optimizer.zero_grad()
    
    out = model(data.x, data.edge_index)
    loss = F.mse_loss(out[data.train_mask].squeeze(), data.y[data.train_mask].squeeze())
    
    loss.backward()
    optimizer.step()

    return loss.item()
    
@torch.no_grad()
def test():
    model.eval()
    
    out = model(data.x, data.edge_index)
    
    train_acc = F.mse_loss(out[train_mask].squeeze(), data.y[train_mask].squeeze())
    test_acc = F.mse_loss(out[test_mask].squeeze(), data.y[test_mask].squeeze())

    return train_acc, test_acc

# Training/Testing
epochs = 200

for epoch in range(0, epochs):
    loss = train()    
    train_acc, test_acc = test()

    if epoch%10 == 0:
        print(f"Epoch: {epoch:03d}, Loss: {loss}, Train: {train_acc:.4f}, Test: {test_acc:.4f}")

# first print is: Epoch: 000, Loss: 57.769744873046875, Train: 54.0927, Test: 112.9581
# last print is: Epoch: 190, Loss: 26.38591957092285, Train: 26.3714, Test: 70.8251

# Inspect output
out = model(data.x, data.edge_index).detach().cpu()
print(out[train_mask].squeeze(), data.y[train_mask].squeeze().cpu())
# produces:
#(tensor([2.3576, 6.4377, 6.4992, 3.1501, 3.1063, 2.8963, 7.0986, 7.0500, 6.7023,
#         3.6353, 6.2670, 5.9594, 6.0270, 2.8551, 3.3796, 2.7211, 7.3768, 6.6157,
#        6.5059, 6.0038, 6.0135, 6.9279, 6.9117, 5.9930, 6.4242]),
#tensor([ 1., 11.,  6.,  7.,  0.,  0., 15., 23., 13., 12.,  9.,  2.,  0.,  0.,
#          3.,  3.,  4.,  2.,  4.,  4.,  2.,  5.,  3.,  2.,  3.]))

print(out[test_mask].squeeze(), data.y[test_mask].squeeze().cpu())
#produces
#(tensor([3.2515, 6.3262, 6.3269, 3.6062, 6.7541, 2.9840, 6.2958, 3.8091, 6.9372,
#        6.0674, 6.3989]),
# tensor([ 9.,  8.,  3., 23., 24.,  5.,  2.,  3.,  3.,  3.,  3.]))

[rusty1s]

Your code looks good to me. Since your labels seems to be only contain integers, you might also want to try out to cast this problem as a classification problem. Furthermore, you may find some success in normalizing targets first. Also keep in mind that you probably do not want to apply a final non-linearity via x = F.relu(x) before returning your targets (depending on whether targets can be negative).

Furthermore, I would try to find out if there are any benefits in applying GraphSAGE at all. For example, you can try to apply linear regression directly on the embedddings obtained by Node2Vec.

[pmoriano]

Hi @rusty1s. I confirm that I tested the two of the ideas you are proposing. (1) Normalizing targets helps indeed. (2) Removing the final non-linearity helps too. I let the GNN definition as below:

# GNN model with 2 layers 
class GNN(torch.nn.Module):
    def __init__(self, in_size, embedding_size, out_size):
        super(GNN, self).__init__()
        self.conv1 = SAGEConv(in_size, embedding_size) # GCNConv, SAGEConv
        self.conv1.aggr = "max"
        self.conv2 = SAGEConv(embedding_size, out_size)
        self.conv2.aggr = "max"

    def forward(self, x, edge_index):
        
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.1, training=self.training)
        
        x = self.conv2(x, edge_index)
        
        return x

I have not tried the linear regression on the node embedding yet. However, I wonder if there is any recommendation to follow with respect to hyperparameter optimization in the GNN setting. There are many different parameters to try including, GNN layer type, embedding size, activation, dropout rate, number of layers, and possibly many more. I am looking for a systematic way to do so that may not be grid search? I would appreciate any pointers/code sample in that regard.

Thanks again.

[rusty1s]

I'm no expert in AutoML, but have you looked at GraphGym? This gives a good overview of the GNN design space. It utilizes random search for hyperparameter tuning which works well in practice.

[pmoriano]

Thanks for the answer @rusty1s. I was not aware of GraphGym but it looks useful for hyperparameter optimization. I am checking it out.

Regarding the question for which I opened this discussion, i.e., improving the performance in the node regression task. I used the code above to train a GNN. After training, the train error is low and the predictions if the training set are good (measured through MSE). However, when I predict on the test dataset, the performance is not good. Note that I am using the whole network structure and node embeddings (node2vec) and have missing labels for training. To me, this looks like a symptom of overfitting. Note that I am using a dropout layer in the GNN. Besides this, could you please offer practical advise on how to improve the prediction on the test set?

Your help is always appreciated. Thanks again.

[rusty1s]

You should be able to measure when overfitting occurs directly during training, e.g., by making use of an additional validation set. You can then evaluate on the model that performs best on the validation set.

[pmoriano]

Thanks @rusty1s. It makes sense. I will consider this closed and let you know about my progress. Your help is highly appreciated.