Model training fits a GNN to prepared graph data by updating its weights. Training, validation, and testing serve distinct roles in the workflow. section 3.1.3
A practical training loop proceeds in stages:
- Initialize an Adam optimizer for each model, using a learning rate of 0.01. section 3.1.3
- For each epoch, run the model on graph data, calculate loss on training examples, backpropagate, and update weights. section 3.1.3
- Validate on unseen data and track training and validation losses to monitor overfitting and learning progress. section 3.1.3
- After training, evaluate on a separate test set to assess generalization. section 3.1.3
Listing 3.4 Training loop
gcn_model = GCN(in_channels=dataset.num_features,\
hidden_channels=64, out_channels=\
dataset.num_classes) #1
graphsage_model = GraphSAGE(in_channels=d\
ataset.num_features, hidden_channels=64, \
out_channels=dataset.num_classes) #2
optimizer_gcn = torch.optim.Adam\
(gcn_model.parameters(), lr=0.01) #3
optimizer_sage = torch.optim.Adam(\
graphsage_model.parameters(), lr=0.01) #4
criterion = torch.nn.CrossEntropyLoss() #5
def train(model, optimizer, data): #6
model.train()
optimizer.zero_grad()
out = model(data.x, data.edge_index)
loss = criterion(out[data.train_mask], data.y[data.train_mask].squeeze())
loss.backward()
optimizer.step()
return loss.item()
def validate(model, data): #7
model.eval()
with torch.no_grad():
out = model(data.x, data.edge_index)
val_loss = criterion(out[data.val_mask], data.y[data.val_mask].squeeze())
return val_loss.item()
train_loss_gcn = [] #8
val_loss_gcn = []
train_loss_sage = []
val_loss_sage = []
for epoch in range(200): #9
loss_gcn = train(gcn_model, optimizer_gcn, subset_graph)
train_loss_gcn.append(loss_gcn)
val_loss_gcn.append(validate(gcn_model, subset_graph))
loss_sage = train(graphsage_model, optimizer_sage, subset_graph)
train_loss_sage.append(loss_sage)
val_loss_sage.append(validate(graphsage_model, subset_graph))
if epoch % 10 == 0:
print(f'Epoch {epoch}, GCN Loss: \
{loss_gcn:.4f}, GraphSAGE Loss: \
{loss_sage:.4f}, GCN Val Loss: \
{val_loss_gcn[-1]:.4f}, GraphSAGE \
Val Loss: {val_loss_sage[-1]:.4f}')
The training-loop listing shows GCN and GraphSAGE models, optimizers, loss calculation, and parameter updates. section 3.1.3
