1
$\begingroup$

I am working on a node classification problem with graphical data. I've used PyTorch to classify nodes by simply applying a network to the individual nodes (e.g., ignoring graphical structure), and I've used PyTorch Geometric to classify nodes by applying a GNN (e.g., GCN).

Is it possible to apply a model from PyTorch as the final layer in PyTorch geometric? I am a bit confused about how this would work because the inputs to the torch.nn.Module are a graph for the GNN case but feature vectors for the MLP case. I want to combine a PyTorch model and a PyTorch Geometric model into a single model, which I can train. Is this possible?

$\endgroup$
2
  • 1
    $\begingroup$ You can define a combined model by creating an instance of both a PyTorch model and a PyTorch Geometric model, and then create a forward pass that applies both models to the input data. $\endgroup$
    – Vic
    Commented Dec 17, 2022 at 5:36
  • $\begingroup$ @Vic rewriting comment because typos. so I create 3 models? 1 in PyTorch, 1 in PyG, then use each of those models as layers in a 3rd model? Will the parameters for each of the first 2 models be updated while training the 3rd model (which includes the others as layers)? $\endgroup$
    – Ralff
    Commented Dec 17, 2022 at 17:31

1 Answer 1

2
$\begingroup$

Continuing your comment, Yes, that's correct! This is just one way of doing it. You can create 3 models: 1 in PyTorch, 1 in PyTorch Geometric, and then use each of those models as layers in a 3rd model.

The parameters for each of the first 2 models will be updated while training the 3rd model, as long as you have defined the 3rd model to contain the first 2 models as submodules and called their forward methods in the forward pass of the 3rd model.

You can refer to the following example:

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch_geometric

# Define the PyTorch model
class MLP(nn.Module):
  def __init__(self, input_size, hidden_size, output_size):
    super(MLP, self).__init__()
    self.fc1 = nn.Linear(input_size, hidden_size)
    self.fc2 = nn.Linear(hidden_size, output_size)
    
  def forward(self, x):
    x = F.relu(self.fc1(x))
    x = self.fc2(x)
    return x

# Define the PyTorch Geometric model
class GNN(torch_geometric.nn.Module):
  def __init__(self, input_size, hidden_size, output_size, mlp):
    super(GNN, self).__init__()
    self.conv1 = torch_geometric.nn.GCNConv(input_size, hidden_size)
    self.conv2 = torch_geometric.nn.GCNConv(hidden_size, output_size)
    self.mlp = mlp
    
  def forward(self, x, edge_index):
    x = self.conv1(x, edge_index)
    x = F.relu(x)
    x = self.conv2(x, edge_index)
    x = self.mlp(x)
    return x

# Define the combined model
class CombinedModel(nn.Module):
  def __init__(self, mlp, gnn):
    super(CombinedModel, self).__init__()
    self.mlp = mlp
    self.gnn = gnn
    
  def forward(self, x, edge_index):
    # Apply the MLP to the feature vectors
    x = self.mlp(x)
    # Apply the GNN to the graph and feature vectors
    x = self.gnn(x, edge_index)
    return x

# Create instances of the PyTorch and PyTorch Geometric models
mlp = MLP(input_size=10, hidden_size=20, output_size=5)
gnn = GNN(input_size=20, hidden_size=30, output_size=10, mlp=mlp)

# Create an instance of the combined model
combined_model = CombinedModel(mlp, gnn)
$\endgroup$
1
  • $\begingroup$ This example is really helpful! Thank you! You included the mlp in both the gnn and combined model. The combined model is equivalent if you replace the mpl layer in the GNN with F.relu and nn.Linear instead of using the mlp twice (once in gnn and once in combined) in the combined model, correct? i.e., the mpl in the GNN and combined model are treated as separate layers with separate paremeters? $\endgroup$
    – Ralff
    Commented Dec 18, 2022 at 5:11

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.