Implementation Optimizing PyTorch Training
This example demonstrates how to implement the discussed optimization techniques for training a simple CNN model on the MNIST handwritten digit classification dataset:
1. Import Necessary Libraries
We are importing the required libraries for PyTorch, data processing, visualization, and profiling.
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms
from torch.utils.tensorboard import SummaryWriter
import torch.profiler as profiler
2. Check GPU Availability and Define Model
Let’s check if GPU is available and define a simple convolutional neural network (CNN) model.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
3. Load and Prepare the Dataset
We are now loading the MNIST dataset, to perform data transformations, and create data loaders for training and validation
# Load MNIST dataset
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
# Define data loader
train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(dataset=val_dataset, batch_size=64, shuffle=False)
4. Instantiate the model
# Instantiate the model, loss function, and optimizer
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
4. Define Training and Validation Functions
Training a machine learning model involves several steps to optimize its parameters for better performance. The process typically starts with setting the model to training mode. Next, the training dataset is divided into batches
- Before computing gradients for the new batch, it’s essential to clear the gradients of all optimized parameters. This is achieved by optimizer.zero_grad().
- The model performs a forward pass outputs = model(inputs) to obtain predictions, followed by computing the loss between the predicted outputs and actual labels using a specified loss function loss = criterion(outputs, labels).
- Once the loss is computed, a backward pass is performed loss.backward() to compute the gradients of the loss with respect to the model parameters. These gradients are then used to update the model parameters using the chosen optimization algorithm optimizer.step().
During validation, a similar process is followed, but without updating the model parameters. The validation dataset is iterated through batches.
# Train the model without optimization strategies
def train(model, train_loader, criterion, optimizer, device):
model.train()
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# Validation function
def validate(model, val_loader, criterion, device):
model.eval()
total_correct = 0
total_samples = 0
with torch.no_grad():
for inputs, labels in val_loader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs, 1)
total_samples += labels.size(0)
total_correct += (predicted == labels).sum().item()
accuracy = total_correct / total_samples
return accuracy
5. Function to Log Results in TensorBoard
# Function to log results in TensorBoard
def log_results(writer, epoch, loss, accuracy):
writer.add_scalar('Loss/train', loss, epoch)
writer.add_scalar('Accuracy/val', accuracy, epoch)
6. Train and Log Results Without Optimizations
# Train and log results without optimizations
with SummaryWriter(log_dir='logs/original') as writer:
for epoch in range(5):
train(model, train_loader, criterion, optimizer, device)
accuracy = validate(model, val_loader, criterion, device)
print(f'Epoch {epoch + 1}, Accuracy: {accuracy}')
log_results(writer, epoch, 0, accuracy)
Output:
Epoch 1, Accuracy: 0.9673333333333334
Epoch 2, Accuracy: 0.9745
Epoch 3, Accuracy: 0.9733333333333334
Epoch 4, Accuracy: 0.9685833333333334
Epoch 5, Accuracy: 0.9748333333333333
Using Optimization Strategies
The Code initialize the training data loader with varying batch sizes and optimization strategies.
- The first two lines define data loaders with a batch size of 64, while the next two lines experiment with a larger batch size of 128 for improved GPU utilization.
- Additionally, torch.cuda.amp.GradScaler() is used to apply automatic mixed precision (AMP) for faster training by scaling the loss to prevent numerical underflow.
- Finally, the model is compiled into a torch script using torch.jit.script() to enable graph mode optimization for improved computational efficiency during training
# Apply optimization strategies
# A. Multi-process Data Loading
# Use multi-process data loading for faster data loading
train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True)
# B. Memory Pinning
# Enable memory pinning for faster data transfer
train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True, pin_memory=True)
# C. Increase Batch Size
# Experiment with a larger batch size for improved GPU utilization
train_loader = DataLoader(dataset=train_dataset, batch_size=128, shuffle=True, pin_memory=True)
# D. Reduce Host to Device Copy
# Use memory pinning and increase batch size to minimize copy overhead
train_loader = DataLoader(dataset=train_dataset, batch_size=128, shuffle=True, pin_memory=True)
# E. Set Gradients to None
# Directly set gradients to None for efficient zeroing of gradients
def zero_grad(model):
for param in model.parameters():
param.grad = None
# F. Automatic Mixed Precision (AMP)
# Utilize automatic mixed precision for faster training
scaler = torch.cuda.amp.GradScaler()
# G. Train in Graph Mode
# Enable torch.jit.graph mode for improved computational efficiency
model = torch.jit.script(model)
Final Results after Optimizations
# The final results after optimizations
with SummaryWriter(log_dir='logs/optimized') as writer:
for epoch in range(5):
model.train()
total_loss = 0
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
# AMP: Scale the loss to prevent underflow
with torch.cuda.amp.autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
# AMP: Unscales the gradients and performs optimization
scaler.step(optimizer)
scaler.update()
total_loss += loss.item()
accuracy = validate(model, val_loader, criterion, device)
print(f'Epoch {epoch + 1}, Loss: {total_loss}, Accuracy: {accuracy}')
log_results(writer, epoch, total_loss, accuracy)
Output:
Epoch 1, Loss: 6.215116824023426, Accuracy: 0.9796666666666667
Epoch 2, Loss: 4.03949194191955, Accuracy: 0.9791666666666666
Epoch 3, Loss: 3.299138018861413, Accuracy: 0.9793333333333333
Epoch 4, Loss: 2.995982698048465, Accuracy: 0.979
Epoch 5, Loss: 2.477495740808081, Accuracy: 0.9796666666666667
Accelerate Your PyTorch Training: A Guide to Optimization Techniques
PyTorch’s flexibility and ease of use make it a popular choice for deep learning. To attain the best possible performance from a model, it’s essential to meticulously explore and apply diverse optimization strategies. This article explores effective methods to enhance the training efficiency and accuracy of your PyTorch models.
Table of Content
- Understanding Performance Challenges
- Optimization Techniques for Faster Training
- 1. Multi-process Data Loading
- 2. Memory Pinning
- 3. Increase Batch Size
- 4. Reduce Host to Device Copy
- 5. Set Gradients to None
- 6. Automatic Mixed Precision (AMP)
- 7. Train in Graph Mode
- Implementation Example: Optimizing a CNN for MNIST Classification
Contact Us