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)