Best model saving pattern in PyTorch - ML Experiment: Train & Evaluate

Experiment - Best model saving pattern

Problem:You have trained a PyTorch model for image classification. The training accuracy is high, but you want to save the best model based on validation accuracy during training to avoid losing the best performing model.

Current Metrics:Training accuracy: 95%, Validation accuracy: 80%, but the saved model is always the last epoch model which has validation accuracy 75%.

Issue:The saved model is not the best one because the saving happens only at the end of training, causing loss of the best validation accuracy model.

Your Task

Implement a model saving pattern that saves the model only when the validation accuracy improves, ensuring the best model is saved.

Use PyTorch framework.

Do not save the model every epoch, only when validation accuracy improves.

Keep the training loop structure simple and clear.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# Dummy dataset
X_train = torch.randn(100, 10)
y_train = torch.randint(0, 2, (100,))
X_val = torch.randn(20, 10)
y_val = torch.randint(0, 2, (20,))

train_ds = TensorDataset(X_train, y_train)
val_ds = TensorDataset(X_val, y_val)

train_loader = DataLoader(train_ds, batch_size=16)
val_loader = DataLoader(val_ds, batch_size=16)

# Simple model
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 2)
    def forward(self, x):
        return self.fc(x)

model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

best_val_acc = 0.0
num_epochs = 10

for epoch in range(num_epochs):
    model.train()
    for xb, yb in train_loader:
        optimizer.zero_grad()
        preds = model(xb)
        loss = criterion(preds, yb)
        loss.backward()
        optimizer.step()

    # Validation
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for xb, yb in val_loader:
            preds = model(xb)
            predicted = preds.argmax(dim=1)
            correct += (predicted == yb).sum().item()
            total += yb.size(0)
    val_acc = correct / total
    print(f"Epoch {epoch+1}, Validation Accuracy: {val_acc:.4f}")

    # Save best model
    if val_acc > best_val_acc:
        best_val_acc = val_acc
        torch.save(model.state_dict(), "best_model.pth")
        print(f"Saved best model with accuracy: {best_val_acc:.4f}")

Added tracking of best validation accuracy with variable best_val_acc.

Added condition to save model weights only if current validation accuracy is higher than best_val_acc.

Used torch.save(model.state_dict(), 'best_model.pth') to save the best model.

Results Interpretation

Before: The model saved was always from the last epoch with validation accuracy 75%, which is lower than the best validation accuracy achieved (80%).

After: The model saved is the one with the highest validation accuracy (80%), ensuring the best performing model is kept.

Saving the model only when validation accuracy improves prevents losing the best model during training and helps in deploying the most accurate model.

Bonus Experiment

Modify the saving pattern to also save the optimizer state along with the model to resume training exactly from the best checkpoint.

💡 Hint

Save a dictionary with model.state_dict() and optimizer.state_dict() using torch.save, and load them together to resume training.