Question about training heterogenous GNN model with PyTorch Lightning (to_hetero changes the type of the model)

[mtoshevska]

Hello!

I have a GNN-based model for heterogeneous graph classification. The model is created with homogenous GNN layers and is then transformed into a heterogenous one with the to_hetero function.

class GraphSAGE(nn.Module):
	def __init__(self, num_classes):
		super(GraphSAGE, self).__init__()
		self.conv1 = SAGEConv((-1, -1), 64, add_self_loops=True, dropout=0.6)
		self.conv2 = SAGEConv((-1, -1), 32, add_self_loops=True, dropout=0.6)
		self.conv3 = SAGEConv((-1, -1), 16, add_self_loops=True, dropout=0.6)
		self.linear = nn.Linear(16, num_classes)

	def forward(self, x, edge_index):
		x = self.conv1(x, edge_index).relu()
		x = self.conv2(x, edge_index).relu()
		x = self.conv3(x, edge_index).relu()
		x = self.linear(x).sigmoid()
		return x

The dataset is an instance of the HeteroData that follows the structure provided here. To enable mini-batch training NeighborLoader instances are created.

train_input_nodes = ('user', dataset['user'].train_mask)
val_input_nodes = ('user', dataset['user'].val_mask)
test_input_nodes = ('user', dataset['user'].test_mask)

train_loader = NeighborLoader(dataset, num_neighbors=[8] * 2, shuffle=True, input_nodes=train_input_nodes, batch_size=16)
val_loader = NeighborLoader(dataset, num_neighbors=[8] * 2, input_nodes=val_input_nodes, batch_size=16)
test_loader = NeighborLoader(dataset, num_neighbors=[8] * 2, input_nodes=test_input_nodes, batch_size=16)

When I train/test the model with a standard PyTorch training loop the process continues without error and the model is able to train. However, using PyTorch Lightning to train/test is not possible because the to_hetero function transforms the model back to GraphModule while the PyTorch Lightning Trainer expects an instance of LightningModule.

How can I train heterogeneous GNN models with PyTorch Lightning?

Thanks in advance!

[wsad1]

We have several examples of using lightning for pyg models here. You might have to wrap you model and use some other interfaces Pyg provides to use pytroch-lightning with pyg.

[mtoshevska]

Hi! I tried wrapping the model as in the script relational_gnn.py, but I keep getting the following error when calling the fit function of the Trainer: "RuntimeError: Can't hash IValues with tag 'String'". From the error traceback, the problem seems to be in the data loader (it is an instance of LightningNodeData as suggested in the sample script).

[akihironitta]

However, using PyTorch Lightning to train/test is not possible because the to_hetero function transforms the model back to GraphModule while the PyTorch Lightning Trainer expects an instance of LightningModule.

@mtoshevska I haven't tried this yet, but simply applying to_hetero to your pure PyTorch model instead of your LightningModule might resolve your issue:

class YourLightningModule(pl.LightningModule):
    def __init__(self, ...):
        ...
        self.model = to_hetero(GraphSAGE(...), ...)

[mtoshevska]

This solves the issue with the to_hetero function but produces a new error for which I have created a new thread.