A training loop repeatedly feeds data through a model, computes loss, updates parameters, and checks whether learning generalizes. For graph models, it can train GCN and GraphSAGE under identical conditions so architectural differences remain comparable. section 3.1.3
A practical loop typically follows these stages:
- Initialize the model, optimizer, and loss function. section 3.1.3
- For each training batch or epoch, switch to training mode, clear old gradients, run a forward pass, calculate loss, backpropagate, and update parameters. section 3.1.3 section 6.4.3
- Run validation without parameter updates, record training and validation losses, and inspect learning curves for overfitting. section 3.1.3
- Finish with test-set evaluation using task-specific metrics or future-prediction error. section 3.1.3 section 6.4.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}')
Inspect the PyTorch implementation of parallel GCN and GraphSAGE training, including model setup, optimization, masking, and validation. section 3.1.3
Listing 6.21 NRI training loop
pbar = tqdm(range(start_epoch, num_epochs + 1), desc='Epochs')
for epoch in pbar:
model.train() #1
model.train_percent = epoch / num_epochs
total_training_loss = 0
for batch in train_data_loader:
inputs = batch['inputs'].cuda(non_blocking=True)
loss, _, _, _, _ = model.\
calculate_loss(inputs,
is_train=True,
return_logits=True)
loss.backward() #2
optimizer.step()
optimizer.zero_grad() #3
total_training_loss += loss.item()
if training_scheduler is not None:
training_scheduler.step()
total_nll, total_kl = 0, 0
for batch in val_data_loader:
inputs = batch['inputs'].cuda(non_blocking=True)
, loss_nll, loss_kl, _, _ = model.calculate_loss(inputs,
is_train=False,
teacher_forcing=True,
return_logits=True)
total_kl += loss_kl.sum().item()
total_nll += loss_nll.sum().item()
total_kl /= len(val_data)
total_nll /= len(val_data)
total_loss = total_kl + total_nll
tuning_loss = total_nll
if tuning_loss < best_val_result:
best_val_epoch, best_val_result = epoch, tuning_loss
Inspect an extended loop with batched NRI training, validation loss components, and a learning-rate scheduler. section 6.4.3
