Num workers for parallel loading in PyTorch

Introduction

Using multiple workers helps load data faster by doing many tasks at the same time. This makes training your model quicker and smoother.

When training a model on a large dataset that takes time to load.

When you want to keep the GPU busy without waiting for data.

When your data loading involves slow operations like reading images from disk.

When you want to speed up the training process by parallelizing data preparation.

Syntax

PyTorch

DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

num_workers sets how many subprocesses load data in parallel.

More workers can speed up loading but use more CPU and memory.

Examples

Loads data in the main process without parallelism. Simple but slower.

PyTorch

DataLoader(dataset, batch_size=64, shuffle=True, num_workers=0)

Uses 2 subprocesses to load data in parallel, speeding up loading.

PyTorch

DataLoader(dataset, batch_size=32, shuffle=True, num_workers=2)

Uses 8 workers without shuffling data, useful for evaluation.

PyTorch

DataLoader(dataset, batch_size=16, shuffle=False, num_workers=8)

Sample Model

This code creates a small dataset and loads it with 4 workers in parallel. It prints how many batches are in one pass over the data.

PyTorch

import torch
from torch.utils.data import DataLoader, TensorDataset

# Create a simple dataset of 100 samples
x = torch.randn(100, 3)
y = torch.randint(0, 2, (100,))
dataset = TensorDataset(x, y)

# DataLoader with 4 workers
loader = DataLoader(dataset, batch_size=10, shuffle=True, num_workers=4)

# Iterate over one epoch and count batches
batch_count = 0
for batch_x, batch_y in loader:
    batch_count += 1
print(f"Total batches: {batch_count}")

Important Notes

Setting num_workers=0 means data loads in the main process, which is slower but simpler.

Too many workers can cause your computer to slow down or run out of memory.

On Windows, use if __name__ == '__main__': guard when using num_workers > 0.

Summary

num_workers controls how many parallel processes load data in PyTorch DataLoader.

Using more workers speeds up data loading but uses more CPU and memory.

Start with a small number like 2 or 4 and increase if your system can handle it.