Num workers for parallel loading in PyTorch - ML Experiment: Train & Evaluate

Experiment - Num workers for parallel loading

Problem:You are training a PyTorch model on image data. The data loading is slow, causing the GPU to wait idle. Currently, the DataLoader uses num_workers=0, which means data loading is done in the main process.

Current Metrics:Training time per epoch: 120 seconds; GPU utilization: 40%; Validation accuracy: 85%

Issue:Data loading is a bottleneck, slowing training and reducing GPU usage efficiency.

Your Task

Increase the number of workers in the DataLoader to speed up data loading and reduce training time per epoch to under 80 seconds, while maintaining validation accuracy above 85%.

Do not change the model architecture or training hyperparameters.

Only modify the DataLoader's num_workers parameter.

Ensure the code runs without errors on a typical multi-core CPU.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import time

# Define dataset and transform
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)

# Create DataLoader with increased num_workers
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)

# Dummy model and optimizer for demonstration
model = torch.nn.Sequential(
    torch.nn.Flatten(),
    torch.nn.Linear(28*28, 128),
    torch.nn.ReLU(),
    torch.nn.Linear(128, 10)
)

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Training loop with timing
model.train()
start_time = time.time()
for images, labels in train_loader:
    images, labels = images.to(device), labels.to(device)
    optimizer.zero_grad()
    outputs = model(images)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
end_time = time.time()

print(f'Training time for one epoch: {end_time - start_time:.2f} seconds')

Increased DataLoader's num_workers from 0 to 4 to enable parallel data loading.

Kept all other code and hyperparameters unchanged.

Results Interpretation

Before: Training time = 120s, GPU utilization = 40%, Validation accuracy = 85%

After: Training time = 75s, GPU utilization = 75%, Validation accuracy = 85%

Increasing num_workers in PyTorch DataLoader allows parallel data loading, reducing training time and improving GPU usage without affecting model accuracy.

Bonus Experiment

Try different num_workers values (e.g., 2, 8, 16) and observe how training time and stability change.

💡 Hint

Too many workers may cause diminishing returns or errors; find the optimal number for your CPU cores.

Practice

1. What does the num_workers parameter control in a PyTorch DataLoader?

easy

A. The number of epochs to train

B. The batch size of the data

C. The number of parallel processes used to load data

D. The learning rate of the model

Solution

1. Step 1: Understand the role of DataLoader

   The DataLoader loads data in batches for training or testing.

2. Step 2: Identify what num_workers does

   num_workers sets how many separate processes load data in parallel to speed up loading.

3. Final Answer:

   The number of parallel processes used to load data -> Option C

4. Quick Check:

   num_workers controls parallel loading = A [OK]

Hint: num_workers means parallel data loading processes [OK]

Common Mistakes:

• Confusing num_workers with batch size
• Thinking num_workers controls training epochs
• Mixing num_workers with learning rate

2. Which of the following is the correct way to set 4 workers in a PyTorch DataLoader?

easy

A. DataLoader(dataset, batch_size=32, workers=4)

B. DataLoader(dataset, batch_size=32, num_workers=4)

C. DataLoader(dataset, batch_size=32, num_worker=4)

D. DataLoader(dataset, batch_size=32, n_workers=4)

Solution

1. Step 1: Recall the exact parameter name

   The correct parameter name is num_workers in PyTorch DataLoader.

2. Step 2: Check the syntax of each option

   Only DataLoader(dataset, batch_size=32, num_workers=4) uses num_workers=4 correctly; others use incorrect parameter names.

3. Final Answer:

   DataLoader(dataset, batch_size=32, num_workers=4) -> Option B

4. Quick Check:

   Correct parameter is num_workers = B [OK]

Hint: Use exact parameter name 'num_workers' in DataLoader [OK]

Common Mistakes:

• Using 'workers' instead of 'num_workers'
• Using singular 'num_worker'
• Using 'n_workers' which is invalid

3. Consider this code snippet:

from torch.utils.data import DataLoader
loader = DataLoader(dataset, batch_size=16, num_workers=2)
for batch in loader:
    print(batch.shape)

What is the effect of setting num_workers=2 here?

medium

A. The model trains with 2 GPUs automatically

B. Only 2 batches will be loaded in total

C. The batch size is effectively doubled

D. Data batches are loaded using 2 parallel worker processes

Solution

1. Step 1: Understand num_workers effect on loading

   Setting num_workers=2 means two separate processes load data in parallel.

2. Step 2: Clarify what does not change

   Batch size remains 16, total batches depend on dataset size, and GPUs are unrelated to this setting.

3. Final Answer:

   Data batches are loaded using 2 parallel worker processes -> Option D

4. Quick Check:

   num_workers=2 means 2 parallel loaders = D [OK]

Hint: num_workers sets parallel loaders, not batch count or GPUs [OK]

Common Mistakes:

• Thinking num_workers limits total batches
• Assuming batch size changes with num_workers
• Confusing num_workers with GPU count

4. You set num_workers=4 in your DataLoader but get an error: RuntimeError: DataLoader worker (pid(s) ...) exited unexpectedly. What is a likely cause?

medium

A. Your system does not have enough CPU cores or memory for 4 workers

B. You forgot to set batch_size

C. You must set num_workers=0 for GPU training

D. The dataset is empty

Solution

1. Step 1: Understand the error cause

   Worker exit errors often happen when system resources (CPU/memory) are insufficient for the number of workers.

2. Step 2: Evaluate other options

   Missing batch_size causes different errors; GPU training does not require num_workers=0; empty dataset causes no batches, not worker crashes.

3. Final Answer:

   Your system does not have enough CPU cores or memory for 4 workers -> Option A

4. Quick Check:

   Worker exit error = resource limits = C [OK]

Hint: Worker errors often mean too many workers for system [OK]

Common Mistakes:

• Assuming batch_size causes worker crashes
• Thinking GPU needs num_workers=0
• Blaming empty dataset for worker exit errors

5. You want to speed up data loading on a machine with 8 CPU cores and 16GB RAM. Which num_workers setting is best to start with?

hard

A. 8 (one worker per CPU core)

B. 2 (small parallel loading)

C. 0 (single process loading)

D. 16 (double the CPU cores)

Solution

1. Step 1: Consider system resources

   With 8 CPU cores, using up to 8 workers can maximize parallel loading without oversubscribing CPUs.

2. Step 2: Avoid too many workers

   Setting 16 workers may cause resource contention and slowdowns; 0 or 2 workers underuse CPU.

3. Final Answer:

   8 (one worker per CPU core) -> Option A

4. Quick Check:

   Match num_workers to CPU cores = A [OK]

Hint: Start num_workers equal to CPU cores for best speed [OK]

Common Mistakes:

• Using 0 workers and missing parallel speedup
• Setting too many workers causing slowdowns
• Ignoring system CPU core count