DataParallel basics in PyTorch - ML Experiment: Train & Evaluate

Experiment - DataParallel basics

Problem:You have a neural network model training on a single GPU, but training is slow. You want to use multiple GPUs to speed up training using PyTorch's DataParallel.

Current Metrics:Training time per epoch: 120 seconds, Validation accuracy: 75%, Training accuracy: 80%

Issue:Training is slow because only one GPU is used. The model does not utilize multiple GPUs to speed up training.

Your Task

Modify the existing PyTorch training code to use DataParallel to run the model on multiple GPUs and reduce training time per epoch by at least 40%, while maintaining validation accuracy above 75%.

You must use PyTorch's DataParallel module.

Do not change the model architecture.

Keep batch size and other hyperparameters the same.

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
from torchvision import datasets, transforms

# Simple CNN model
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(1, 10, kernel_size=5)
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(10 * 12 * 12, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv(x)))
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# Data preparation
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST('.', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

val_dataset = datasets.MNIST('.', train=False, download=True, transform=transform)
val_loader = DataLoader(val_dataset, batch_size=1000, shuffle=False)

# Device setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Model setup
model = SimpleCNN()
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)
model.to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())

# Training loop
for epoch in range(1):  # single epoch for demo
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    for inputs, targets in train_loader:
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * inputs.size(0)
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

    train_loss = running_loss / total
    train_acc = 100. * correct / total

    # Validation
    model.eval()
    val_loss = 0.0
    val_correct = 0
    val_total = 0
    with torch.no_grad():
        for inputs, targets in val_loader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            val_loss += loss.item() * inputs.size(0)
            _, predicted = outputs.max(1)
            val_total += targets.size(0)
            val_correct += predicted.eq(targets).sum().item()

    val_loss /= val_total
    val_acc = 100. * val_correct / val_total

    print(f'Epoch {epoch+1}: Train Loss={train_loss:.4f}, Train Acc={train_acc:.2f}%, Val Loss={val_loss:.4f}, Val Acc={val_acc:.2f}%')

Wrapped the model with nn.DataParallel to enable multi-GPU training.

Moved the model to the CUDA device after wrapping with DataParallel.

Ensured inputs and targets are moved to the correct device inside the training loop.

Results Interpretation

Before: Training time per epoch: 120s, Validation accuracy: 75%, Training accuracy: 80%

After: Training time per epoch: 70s, Validation accuracy: 76%, Training accuracy: 81%

Using DataParallel allows the model to use multiple GPUs, which speeds up training significantly without hurting accuracy. This shows how parallel computing can improve efficiency in deep learning.

Bonus Experiment

Try using torch.nn.parallel.DistributedDataParallel instead of DataParallel for potentially better multi-GPU performance.

💡 Hint

DistributedDataParallel requires setting up a process group and launching multiple processes, but it can reduce overhead compared to DataParallel.

Practice

1. What is the main purpose of torch.nn.DataParallel in PyTorch?

easy

A. To convert a model to run only on CPU

B. To run a model on multiple GPUs by splitting input data automatically

C. To reduce the size of the model parameters

D. To change the model architecture dynamically during training

Solution

1. Step 1: Understand DataParallel's role

   DataParallel helps use multiple GPUs by splitting input data and running model copies in parallel.

2. Step 2: Compare options

   Only To run a model on multiple GPUs by splitting input data automatically describes this behavior correctly; others describe unrelated tasks.

3. Final Answer:

   To run a model on multiple GPUs by splitting input data automatically -> Option B

4. Quick Check:

   DataParallel = multi-GPU input splitting [OK]

Hint: DataParallel splits input across GPUs automatically [OK]

Common Mistakes:

• Thinking DataParallel changes model architecture
• Confusing DataParallel with CPU-only mode
• Believing it reduces model size

2. Which of the following is the correct way to wrap a PyTorch model model with DataParallel?

easy

A. model = torch.DataParallel(model)

B. model = torch.nn.parallel.DataParallel(model)

C. model = torch.nn.DataParallel(model)

D. model = torch.parallel.DataParallel(model)

Solution

1. Step 1: Recall correct import path

   DataParallel is in torch.nn, so correct usage is torch.nn.DataParallel.

2. Step 2: Check options

   model = torch.nn.DataParallel(model) matches the correct usage; others are incorrect or incomplete.

3. Final Answer:

   model = torch.nn.DataParallel(model) -> Option C

4. Quick Check:

   Correct import path = torch.nn.DataParallel [OK]

Hint: DataParallel is in torch.nn module [OK]

Common Mistakes:

• Using torch.DataParallel instead of torch.nn.DataParallel
• Confusing torch.nn.parallel with torch.nn
• Confusing with other torch submodules

3. Given the code below, what will be the output shape of output if input shape is (8, 3, 32, 32) and 2 GPUs are used?

import torch
import torch.nn as nn

model = nn.Linear(32*32*3, 10)
model = nn.parallel.DataParallel(model)
input = torch.randn(8, 3, 32, 32)
input = input.view(8, -1)
output = model(input)
print(output.shape)

medium

A. torch.Size([8, 10])

B. torch.Size([4, 10])

C. torch.Size([16, 10])

D. RuntimeError due to shape mismatch

Solution

1. Step 1: Understand input and model

   Input is reshaped to (8, 3072) and model outputs 10 features per input.

2. Step 2: DataParallel splits batch across 2 GPUs

   Batch of 8 splits into 2 batches of 4, processed in parallel, outputs combined back to (8, 10).

3. Final Answer:

   torch.Size([8, 10]) -> Option A

4. Quick Check:

   Output shape = batch size x output features = (8, 10) [OK]

Hint: Output batch size stays same after DataParallel [OK]

Common Mistakes:

• Thinking output batch size halves with GPUs
• Confusing input reshape with output shape
• Expecting runtime error due to DataParallel

4. What is the likely cause of this error when using DataParallel?

RuntimeError: Expected tensor for argument #1 'input' to have scalar type Float but got Double

Code snippet:

model = torch.nn.parallel.DataParallel(model)
input = torch.randn(8, 3, 32, 32, dtype=torch.double)
output = model(input)

medium

A. Input batch size is too large for GPUs

B. DataParallel requires input to be on CPU only

C. Model is not wrapped correctly with DataParallel

D. Input tensor dtype does not match model parameters dtype

Solution

1. Step 1: Analyze error message

   Error says input tensor is Double but model expects Float tensors.

2. Step 2: Check input and model dtypes

   Model parameters default to Float; input is Double, causing mismatch and error.

3. Final Answer:

   Input tensor dtype does not match model parameters dtype -> Option D

4. Quick Check:

   Dtype mismatch causes runtime error [OK]

Hint: Ensure input and model dtypes match before DataParallel [OK]

Common Mistakes:

• Assuming DataParallel only works on CPU
• Not checking tensor data types
• Thinking batch size causes dtype errors

5. You want to speed up training on 4 GPUs using DataParallel. Your model has a batch normalization layer. What should you be careful about to avoid incorrect training behavior?

hard

A. BatchNorm layers behave differently in DataParallel; use SyncBatchNorm or move to DistributedDataParallel

B. DataParallel automatically fixes BatchNorm issues, no action needed

C. Remove BatchNorm layers before using DataParallel

D. Increase batch size to avoid BatchNorm problems

Solution

1. Step 1: Understand BatchNorm behavior in DataParallel

   BatchNorm computes stats per device, which can cause inconsistent normalization across GPUs.

2. Step 2: Recommended solutions

   Use SyncBatchNorm to sync stats or switch to DistributedDataParallel for correct BatchNorm behavior.

3. Final Answer:

   BatchNorm layers behave differently in DataParallel; use SyncBatchNorm or move to DistributedDataParallel -> Option A

4. Quick Check:

   BatchNorm needs syncing or DDP for multi-GPU [OK]

Hint: BatchNorm needs syncing or DDP on multiple GPUs [OK]

Common Mistakes:

• Assuming DataParallel handles BatchNorm correctly by default
• Removing BatchNorm layers unnecessarily
• Only increasing batch size without syncing stats