Reproducibility and Training State Management in PyTorch: Best Practices and Tips | by Sedjro Salomon Hotegni | Jul, 2024

Within the ever-evolving area of machine studying, reproducibility, and environment friendly coaching state administration are essential for analysis and sensible functions. Whether or not you’re creating new fashions, experimenting with hyperparameters, or deploying options to manufacturing, making certain that your outcomes are constant and that you could seamlessly resume coaching after interruptions can prevent precious time and assets.

On this weblog put up, I’ll present hands-on code examples and detailed explanations that will help you obtain reproducibility in your machine-learning experiments. I’ll additionally share some efficient strategies for managing coaching state in PyTorch, permitting you to renew your long-running coaching processes effortlessly, even after sudden shutdowns or deliberate interruptions. By the top of this text, you’ll have a stable understanding of methods to make your PyTorch workflows extra strong and environment friendly, making certain that your experiments are dependable, and your coaching processes are resilient. Let’s get began!

“Is there a reproducibility disaster in science? Most agree that there’s a disaster and over 70% mentioned they’d tried and failed to breed one other group’s experiments”​ — (Nature)​.

To make sure the reproducibility of your machine studying experiments, it’s essential to set a constant random seed throughout all libraries and frameworks you’re utilizing. The next perform, `set_seed`, accomplishes this in PyTorch by setting seeds for NumPy, Python’s built-in “random” module, and varied PyTorch elements.

def set_seed(seed):
"""Units the seed for reproducibility throughout varied libraries.
:param seed: The seed worth to make sure reproducibility
:return: None
"""
np.random.seed(seed)
random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

• `np.random.seed(seed)`: Units the seed for NumPy’s random quantity generator, making certain any random numbers generated by NumPy are reproducible.
• `random.seed(seed)`: Units the seed for Python’s built-in “random” module, making certain reproducibility for any random numbers generated utilizing this module.
• `torch.manual_seed(seed)`: Units the seed for PyTorch’s random quantity generator, making certain any random numbers generated by PyTorch (for CPU operations) are reproducible.
• `torch.cuda.manual_seed(seed)` and `torch.cuda.manual_seed_all(seed)`: If CUDA is out there, these capabilities set the seed for all GPUs, making certain reproducibility for GPU operations.
• CuDNN settings:
  — `torch.backends.cudnn.enabled = True`: Ensures CuDNN (NVIDIA’s deep neural community library) is enabled for higher efficiency.
  — `torch.backends.cudnn.benchmark = False`: Disables the CuDNN auto-tuner that selects the most effective algorithms on your {hardware}, which might introduce non-deterministic conduct.
  — `torch.backends.cudnn.deterministic = True`: Forces CuDNN to make use of solely deterministic algorithms, making certain that operations produce the identical outcome each time.

Utilizing this perform originally of your script will provide help to obtain constant and reproducible outcomes throughout completely different runs of your machine-learning experiments.

Kohn’s Second Regulation humorously factors out a important side of scientific reproducibility:

“An experiment is reproducible till one other laboratory tries to repeat it.” — (Kohn’s Second Regulation)

This witty commentary underscores the sensible challenges of making certain reproducibility in scientific analysis. In machine studying, reproducibility goes past simply operating the identical code. It additionally entails sustaining constant {hardware} necessities and software program environments throughout all runs.

For instance, the particular kind of GPU or CPU used can have an effect on efficiency and outcomes. Moreover, variations in library variations can result in completely different outcomes, making it important to doc and handle these dependencies. When sharing your work, it’s essential to specify the {hardware} configurations and software program variations used. This manner, others making an attempt to breed your outcomes can guarantee their atmosphere matches yours as carefully as doable.

“The best glory in dwelling lies not in by no means falling, however in rising each time we fall.” — (Nelson Mandela)

To reduce the impression of sudden shutdowns or deliberate interruptions, it can save you the mannequin’s state after every epoch (every new epoch’s checkpoint will overwrite the earlier one). This lets you resume coaching from the final saved state in case of interruptions. Beneath are two important capabilities for saving and loading the coaching state in PyTorch.

Save Checkpoint Perform

The `save_checkpoint` perform saves the present state of the mannequin, optimizer, and studying fee scheduler to a file. This ensures that you could resume coaching from the identical level at a later time.

def save_checkpoint(
epoch,
mannequin,
optimizer,
lr_scheduler,
TRAIN_LOSS,
VAL_ACCU,
main_dir,
mod_logdir,
best_val_accu_at=None, # that is the checklist [best_val_accu, best_epoch]
verbose=True,
):
if not os.path.exists("%s/%s" % (main_dir, mod_logdir)):
os.makedirs("%s/%s" % (main_dir, mod_logdir))
filename = str("%s/%s/last_train_checkpoint.pth" % (main_dir, mod_logdir))
checkpoint = {
"epoch": epoch,
"model_state_dict": mannequin.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"lr_scheduler_state_dict": lr_scheduler.state_dict(),
"last_best_valid_accuracy_epoch": best_val_accu_at,
"list_train_loss": TRAIN_LOSS,
"list_val_accuracy": VAL_ACCU,
}
torch.save(checkpoint, filename)
if verbose:
print(f"Checkpoint saved at {filename}.")

• Listing Creation: The perform checks if the required listing exists and creates it if essential.
• Checkpoint Dictionary: This dictionary consists of the present epoch, mannequin state, optimizer state, studying fee scheduler state, the most effective validation accuracy and corresponding epoch, coaching losses, and validation accuracies.
• Saving Checkpoint: The `torch.save` perform is used to avoid wasting the checkpoint to a file.
• Verbose Logging: If `verbose` is about to `True`, a message indicating the save location is printed.

Load Checkpoint Perform

The `load_checkpoint` perform restores the mannequin, optimizer, and studying fee scheduler states from a beforehand saved checkpoint. This allows you to proceed coaching from the final saved state.

def load_checkpoint(
mannequin,
optimizer,
lr_scheduler,
main_dir,
mod_logdir,
verbose=True,
):
filename = str("%s/%s/last_train_checkpoint.pth" % (main_dir, mod_logdir))
checkpoint = torch.load(filename)
mannequin.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
if lr_scheduler is just not None:
lr_scheduler.load_state_dict(checkpoint["lr_scheduler_state_dict"])
epoch = checkpoint["epoch"]
last_best_val_accu_epoch = checkpoint["last_best_valid_accuracy_epoch"]
list_train_loss = checkpoint["list_train_loss"]
list_val_accuracy = checkpoint["list_val_accuracy"]
if verbose:
print(f"Checkpoint loaded.")
return (
epoch,
last_best_val_accu_epoch,
list_train_loss,
list_val_accuracy,
)

• Loading Checkpoint: The perform hundreds the checkpoint utilizing `torch.load`.
• Restoring States: It restores the states of the mannequin, optimizer, and, if used, the educational fee scheduler from the checkpoint.
• Returning State Info: The perform returns the epoch, finest validation accuracy and epoch, coaching losses, and validation accuracies. No have to return the mannequin, optimizer, and studying fee scheduler, because the loading operation is carried out in place.
• Verbose Logging: If `verbose` is about to `True`, a message indicating that the checkpoint was loaded is printed.

This ensures that you could save the coaching state after every epoch and resume coaching seamlessly, minimizing the impression of interruptions.

On this part, I’ll present an entire instance of coaching a easy Convolutional Neural Community (CNN) to categorise the CIFAR-10 dataset. We’ll use the capabilities we mentioned earlier for reproducibility and coaching state administration. Moreover, we’ll combine Weights and Biases (WandB) to trace and visualize varied points of the coaching course of in real-time (If you’re having hassle connecting to WandB, you’ll be able to skip/remark out the associated code, or search for extra assistance on the official website.).

Step-by-Step Information:

1. Set Up Setting:
— Set up the required libraries.
— Import essential packages.

2. Outline Reproducibility and Checkpoint Capabilities:
— Use the `set_seed`, `save_checkpoint`, and `load_checkpoint` capabilities.

3. Outline the Mannequin:
— Create a easy CNN.

4. Coaching Loop with WandB Integration:
— Prepare the mannequin and log metrics to WandB.

Code:

First, join a free W&B account. Second, set up the W&B SDK with pip. Navigate to your terminal and kind the next command:

pip set up wandb

Coaching:

import os
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.useful as F
import torch.optim as optim
import torch.backends.cudnn as cudnn
import torchvision
import torchvision.transforms as transforms
import wandb
# Step 1: Arrange atmosphere and initialize WandB
wandb.init(mission="cifar10-classification")
# Set reproducibility seed
def set_seed(seed):
"""Units the seed for reproducibility throughout varied libraries.
:param seed: The seed worth to make sure reproducibility
:return: None
"""
np.random.seed(seed)
random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
set_seed(42)
# Step 2: Outline reproducibility and checkpoint capabilities
def save_checkpoint(
epoch,
mannequin,
optimizer,
lr_scheduler,
TRAIN_LOSS,
VAL_ACCU,
main_dir,
mod_logdir,
best_val_accu_at=None, # that is the checklist [best_val_accu, best_epoch]
verbose=True,
):
if not os.path.exists("%s/%s" % (main_dir, mod_logdir)):
os.makedirs("%s/%s" % (main_dir, mod_logdir))
filename = str("%s/%s/last_train_checkpoint.pth" % (main_dir, mod_logdir))
checkpoint = {
"epoch": epoch,
"model_state_dict": mannequin.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"lr_scheduler_state_dict": lr_scheduler.state_dict(),
"last_best_valid_accuracy_epoch": best_val_accu_at,
"list_train_loss": TRAIN_LOSS,
"list_val_accuracy": VAL_ACCU,
}
torch.save(checkpoint, filename)
if verbose:
print(f"Checkpoint saved at {filename}.")
def load_checkpoint(
mannequin,
optimizer,
lr_scheduler,
main_dir,
mod_logdir,
verbose=True,
):
filename = str("%s/%s/last_train_checkpoint.pth" % (main_dir, mod_logdir))
checkpoint = torch.load(filename)
mannequin.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
if lr_scheduler is just not None:
lr_scheduler.load_state_dict(checkpoint["lr_scheduler_state_dict"])
epoch = checkpoint["epoch"]
last_best_val_accu_epoch = checkpoint["last_best_valid_accuracy_epoch"]
list_train_loss = checkpoint["list_train_loss"]
list_val_accuracy = checkpoint["list_val_accuracy"]
if verbose:
print(f"Checkpoint loaded.")
return (
epoch,
last_best_val_accu_epoch,
list_train_loss,
list_val_accuracy,
)
# Step 3: Outline the mannequin
class SimpleCNN(nn.Module):
def __init__(self):
tremendous(SimpleCNN, self).__init__()
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
self.fc1 = nn.Linear(64 * 8 * 8, 512)
self.fc2 = nn.Linear(512, 10)
def ahead(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 64 * 8 * 8)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
# Step 4: Coaching loop with WandB integration
def train_and_evaluate(resume_training = False):
# Hyperparameters
batch_size = 64
learning_rate = 0.001
num_epochs = 10
main_dir = "checkpoints"
mod_logdir = "cifar10_cnn"
# Get gadget
gadget = torch.gadget("cuda:0" if torch.cuda.is_available() else "cpu")
# Information loaders
rework = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
trainset = torchvision.datasets.CIFAR10(root='./knowledge', practice=True, obtain=True, rework=rework)
trainloader = torch.utils.knowledge.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./knowledge', practice=False, obtain=True, rework=rework)
testloader = torch.utils.knowledge.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)
# Mannequin, optimizer, and scheduler
mannequin = SimpleCNN().to(gadget)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(mannequin.parameters(), lr=learning_rate)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
# WandB config
wandb.config.replace({
"batch_size": batch_size,
"learning_rate": learning_rate,
"num_epochs": num_epochs,
})
init_epoch = 0
best_val_accu = 0.0
best_epoch = 0
train_losses = []
val_accuracies = []
if resume_training:
(
last_epoch,
last_best_val_accu_epoch,
TRAIN_LOSS,
VAL_ACCU,
) = load_checkpoint(
mannequin,
optimizer,
lr_scheduler,
main_dir,
mod_logdir,
)
init_epoch = last_epoch + 1
print(f"*** Resuming coaching from epoch {init_epoch + 1} ***")
if last_best_val_accu_epoch is just not None:
best_val_accu, best_epoch = last_best_val_accu_epoch
print(
f"Finest Validation Accuracy to date: {best_val_accu} at epoch {best_epoch + 1}."
)
best_val_accu = best_val_accu
best_epoch = best_epoch
for epoch in vary(init_epoch, num_epochs):
mannequin.practice()
running_loss = 0.0
for inputs, labels in trainloader:
inputs, labels = inputs.to(gadget), labels.to(gadget)
optimizer.zero_grad()
outputs = mannequin(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.merchandise()
train_loss = running_loss / len(trainloader)
train_losses.append(train_loss)
# Validation
mannequin.eval()
right = 0
whole = 0
with torch.no_grad():
for inputs, labels in testloader:
inputs, labels = inputs.to(gadget), labels.to(gadget)
outputs = mannequin(inputs)
_, predicted = torch.max(outputs.knowledge, 1)
whole += labels.measurement(0)
right += (predicted == labels).sum().merchandise()
val_accu = right / whole
val_accuracies.append(val_accu)
# Log metrics to WandB
wandb.log({
"epoch": epoch,
"train_loss": train_loss,
"val_accuracy": val_accu,
})
# Replace finest accuracy
if val_accu > best_val_accu:
best_val_accu = val_accu
best_epoch = epoch
lr_scheduler.step()
# Save checkpoint
save_checkpoint(
epoch,
mannequin,
optimizer,
lr_scheduler,
train_losses,
val_accuracies,
main_dir,
mod_logdir,
best_val_accu_at=[best_val_accu, best_epoch],
)
print(f"Epoch {epoch+1}/{num_epochs}, Prepare Loss: {train_loss:.4f}, Val Accuracy: {val_accu:.4f}")
print(f"Finest Validation Accuracy: {best_val_accu:.4f} at epoch {best_epoch + 1}")

Code Clarification

1. Setting Up Setting:
— Initializes WandB and units the random seed for reproducibility.
— Configures the info loaders for the CIFAR-10 dataset.

2. Save and Load Checkpoint Capabilities:
— `save_checkpoint`: Saves the present coaching state, together with mannequin parameters, optimizer state, studying fee scheduler state, and coaching/validation metrics.
— `load_checkpoint`: Hundreds the saved state to renew coaching.

3. Outline the Mannequin:
— A easy CNN is outlined with two convolutional layers, adopted by absolutely related layers.

4. Coaching Loop:
— The coaching loop runs for a specified variety of epochs, computing coaching loss, and validation accuracy.
— Saves checkpoints after every epoch and logs metrics to WandB for real-time monitoring and visualization.

Right here, I’m utilizing Google Colab with a T4 GPU. Operating the earlier code ought to give an output like this:

wandb: Logging into wandb.ai. (Discover ways to deploy a W&B server regionally: https://wandb.me/wandb-server)
wandb: You will discover your API key in your browser right here: https://wandb.ai/authorize
wandb: Paste an API key out of your profile and hit enter, or press ctrl+c to give up: ··········
wandb: Appending key for api.wandb.ai to your netrc file: /root/.netrc
Monitoring run with wandb model 0.17.4
Run knowledge is saved regionally in /content material/wandb/run-20240713_173200-rlp809er
Syncing run super-snowflake-8 to Weights & Biases (docs)
View mission at [YOUR_PROJECT_LINK]
View run at [LINK_FOR_REAL_TIME_TRACKING]

For those who run the next code with resume_training = False, the coaching begins from scratch:

train_and_evaluate(resume_training = False)

Your output ought to seem like this:

Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ./knowledge/cifar-10-python.tar.gz
100%|██████████| 170498071/170498071 [00:12<00:00, 13141203.42it/s]
Extracting ./knowledge/cifar-10-python.tar.gz to ./knowledge
Information already downloaded and verified
Checkpoint saved at checkpoints/cifar10_cnn/last_train_checkpoint.pth.
Epoch 1/10, Prepare Loss: 1.3100, Val Accuracy: 0.6178
Checkpoint saved at checkpoints/cifar10_cnn/last_train_checkpoint.pth.
Epoch 2/10, Prepare Loss: 0.9143, Val Accuracy: 0.6917
Checkpoint saved at checkpoints/cifar10_cnn/last_train_checkpoint.pth.
Epoch 3/10, Prepare Loss: 0.7298, Val Accuracy: 0.7123
Checkpoint saved at checkpoints/cifar10_cnn/last_train_checkpoint.pth.
Epoch 4/10, Prepare Loss: 0.5696, Val Accuracy: 0.7334
Checkpoint saved at checkpoints/cifar10_cnn/last_train_checkpoint.pth.
Epoch 5/10, Prepare Loss: 0.4237, Val Accuracy: 0.7341

5. Resuming Coaching:
In case your coaching has been interrupted, you’ll be able to resume coaching by setting the resume_training argument to `True`.

train_and_evaluate(resume_training = True)

Output:

Information already downloaded and verified
Information already downloaded and verified
Checkpoint loaded.
*** Resuming coaching from epoch 6 ***
Finest Validation Accuracy to date: 0.7341 at epoch 5.
Checkpoint saved at checkpoints/cifar10_cnn/last_train_checkpoint.pth.
Epoch 6/10, Prepare Loss: 0.2938, Val Accuracy: 0.7274
Checkpoint saved at checkpoints/cifar10_cnn/last_train_checkpoint.pth.
Epoch 7/10, Prepare Loss: 0.1906, Val Accuracy: 0.7290
Checkpoint saved at checkpoints/cifar10_cnn/last_train_checkpoint.pth.
Epoch 8/10, Prepare Loss: 0.0631, Val Accuracy: 0.7477
Checkpoint saved at checkpoints/cifar10_cnn/last_train_checkpoint.pth.
Epoch 9/10, Prepare Loss: 0.0393, Val Accuracy: 0.7483
Checkpoint saved at checkpoints/cifar10_cnn/last_train_checkpoint.pth.
Epoch 10/10, Prepare Loss: 0.0287, Val Accuracy: 0.7491
Finest Validation Accuracy: 0.7491 at epoch 10

On this instance, we guarantee reproducibility, handle coaching states effectively, and visualize the coaching course of with WandB, making our experiment extra dependable and insightful.

Guaranteeing reproducibility and successfully managing the coaching state of your machine studying fashions are essential steps in conducting strong and dependable analysis. By setting a constant random seed, saving checkpoints, and utilizing instruments like Weights and Biases (WandB), you can also make your coaching course of resilient to interruptions and simpler to trace and debug. On this weblog put up, we walked by the important strategies for reaching reproducibility with PyTorch, together with an in depth instance of coaching a CNN on the CIFAR-10 dataset. The supplied code examples not solely exhibit methods to keep reproducibility but additionally methods to handle and resume the coaching state seamlessly. Moreover, integrating WandB permits you to monitor your experiments in real-time, offering precious insights into your coaching course of. By implementing these practices, you’ll be able to improve the reliability and transparency of your machine-learning workflows, finally contributing to extra credible and reproducible analysis. Completely satisfied coding, and will your fashions at all times practice easily!

1- Is there a reproducibility disaster in science? — https://www.nature.com/articles/d41586-019-00067-3

2- WandB — https://wandb.ai/site