Is it possible to combine models in pytorch and pytorch geometric?

Question

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?

Answer 1

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)