Pytorch geometric label prediction using GNN - torch

Question

I'm currently trying to find a way how to get a single label prediction from my GNN.I'd like to create a list of ground truths compared to how the model predicts the label.
My code is as following, ran in google collab:
Should I create a validation step or is there another way about this? I'm very new to this.
# Load the dataset
dataset = TUDataset(
root=".", name="Mutagenicity",
transform=AddSymbols(),
).shuffle()
from torch_geometric.data import DataLoader
loader = DataLoader(dataset, batch_size=32, shuffle=True)
data_batch = next(iter(loader))
print(data_batch)
print("Number of graphs in the batch:",data_batch.num_graphs)
import torch
from torch_scatter import scatter
from torch_geometric.utils import to_dense_adj
class GraphConvolution(torch.nn.Module):
def __init__(self, num_input_features, num_output_features):
super().__init__()
self.W1 = torch.nn.Parameter(torch.randn(num_input_features, num_output_features)*0.1)
self.W2 = torch.nn.Parameter(torch.randn(num_input_features, num_output_features)*0.1)
def forward(self, x, edge_index):
adj = to_dense_adj(edge_index)
neighbors_aggregation = adj # x
out = x # self.W1 + neighbors_aggregation # self.W2
return out
from torch.nn import Linear, LogSoftmax
import torch.nn.functional as F
from torch_geometric.nn import GraphConv, global_add_pool
class GNN(torch.nn.Module):
def __init__(
self,
num_classes,
hidden_dim,
node_features_dim,
edge_features_dim=None
):
super(GNN, self).__init__()
self.hidden_dim = hidden_dim
self.conv1 = GraphConv(node_features_dim, hidden_dim)
self.conv2 = GraphConv(hidden_dim, hidden_dim)
self.conv3 = GraphConv(hidden_dim, hidden_dim)
self.conv4 = GraphConv(hidden_dim, hidden_dim)
self.conv5 = GraphConv(hidden_dim, hidden_dim)
self.fc1 = Linear(hidden_dim, hidden_dim)
self.fc2 = Linear(hidden_dim, num_classes)
self.readout = LogSoftmax(dim=-1)
def forward(self, x, edge_index, batch):
x = F.relu(self.conv1(x, edge_index))
x = F.relu(self.conv2(x, edge_index))
x = F.relu(self.conv3(x, edge_index))
x = F.relu(self.conv4(x, edge_index))
x = F.relu(self.conv5(x, edge_index))
x = global_add_pool(x, batch)
x = F.relu(self.fc1(x))
x = F.dropout(x, p=0.5, training=self.training)
x = self.fc2(x)
return self.readout(x)
from torch_geometric.loader import DataLoader
from tqdm.auto import tqdm
# If possible, we use a GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device:", device)
idx_train_end = int(len(dataset) * .5)
idx_valid_end = int(len(dataset) * .7)
BATCH_SIZE = 128
BATCH_SIZE_TEST = len(dataset) - idx_valid_end
# In the test loader we set the natch size to be equal to the size of the whole test set
loader_train = DataLoader(dataset[:idx_train_end], batch_size=BATCH_SIZE, shuffle=True)
loader_valid = DataLoader(dataset[idx_train_end:idx_valid_end], batch_size=BATCH_SIZE, shuffle=True)
loader_test = DataLoader(dataset[idx_valid_end:], batch_size=BATCH_SIZE_TEST, shuffle=False)
from src import compute_accuracy
def train_model(loader_train, loader_valid, model, optimizer, loss_function, N_EPOCHS):
# Prepare empy lists for logging
train_losses = []
train_accs = []
val_accs = []
for epoch in tqdm(range(N_EPOCHS)):
epoch_loss = 0
for batch in tqdm(loader_train, leave=False):
batch.to(device)
out = model(batch.x, batch.edge_index, batch.batch)
loss = loss_function(out, batch.y.flatten())
optimizer.zero_grad()
loss.backward()
optimizer.step()
acc_train = compute_accuracy(model, loader_train)
acc_valid = compute_accuracy(model, loader_valid)
with torch.no_grad():
train_accs.append(acc_train)
val_accs.append(acc_valid)
train_losses.append(loss)
print(f"Epoch: {epoch}, Loss: {loss}")