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

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?

• 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.
– Vic
Commented Dec 17, 2022 at 5:36
• @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)? Commented Dec 17, 2022 at 17:31

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)

• 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? Commented Dec 18, 2022 at 5:11