PyTorch DataLoader

The DataLoader takes a Dataset and serves its samples to your training loop in batches. It handles batching, shuffling, and parallel data loading automatically. Without a DataLoader, you would manually slice tensors into batches, shuffle indices, and coordinate multiple CPU workers — all tedious work that DataLoader eliminates.

Basic DataLoader Setup

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

X = torch.rand(500, 10)
y = torch.randint(0, 3, (500,))

dataset = TensorDataset(X, y)

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

print(f"Total batches per epoch: {len(loader)}")   # 500 / 32 = ~16 batches

Key Parameters

batch_size:
  Number of samples per batch.
  Common values: 16, 32, 64, 128, 256.
  Larger = faster training, more GPU memory needed.

shuffle:
  True → randomly reorder samples before each epoch.
  Always True for training. False for validation/test.

num_workers:
  Number of CPU worker processes to load data in parallel.
  0 = load on main process (safe, slower).
  2 or 4 = faster loading while GPU trains.

drop_last:
  True → discard the final incomplete batch if dataset
  size is not divisible by batch_size.
  Useful when batch norm requires consistent batch sizes.

Iterating Over a DataLoader

for batch_X, batch_y in loader:
    print(batch_X.shape)   # torch.Size([32, 10])
    print(batch_y.shape)   # torch.Size([32])
    break   # just show first batch

DataLoader in a Training Loop

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader

X = torch.rand(400, 8)
y = torch.rand(400, 1)

dataset = TensorDataset(X, y)
loader  = DataLoader(dataset, batch_size=32, shuffle=True)

model     = nn.Linear(8, 1)
optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn   = nn.MSELoss()

for epoch in range(10):
    model.train()
    epoch_loss = 0.0

    for batch_X, batch_y in loader:
        optimizer.zero_grad()
        pred = model(batch_X)
        loss = loss_fn(pred, batch_y)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()

    print(f"Epoch {epoch+1}  Avg Loss: {epoch_loss / len(loader):.4f}")

Batch Flow Diagram

Dataset: 400 samples
DataLoader batch_size=32

Epoch 1:
  Batch 1:  samples  0-31   → shape [32, 8]
  Batch 2:  samples 32-63   → shape [32, 8]
  ...
  Batch 12: samples 352-383 → shape [32, 8]
  Batch 13: samples 384-399 → shape [16, 8]  ← last batch smaller
                                              (use drop_last=True to skip)

Separate Train and Validation Loaders

from torch.utils.data import random_split

train_set, val_set = random_split(dataset, [320, 80])

train_loader = DataLoader(train_set, batch_size=32, shuffle=True)
val_loader   = DataLoader(val_set,   batch_size=32, shuffle=False)

Pinning Memory for Faster GPU Transfer

When training on GPU, set pin_memory=True. This allocates data in page-locked CPU memory, making the transfer from CPU to GPU significantly faster.

loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,
    pin_memory=True   # GPU speedup
)

Custom collate_fn

By default, DataLoader stacks samples from the batch into tensors. When samples have variable lengths — like sentences of different word counts — you need a custom collate_fn to pad them to equal length before stacking.

def my_collate(batch):
    # batch is a list of (sample, label) tuples
    max_len = max(x[0].size(0) for x in batch)
    padded  = torch.zeros(len(batch), max_len)
    labels  = torch.zeros(len(batch), dtype=torch.long)
    for i, (x, y) in enumerate(batch):
        padded[i, :x.size(0)] = x
        labels[i] = y
    return padded, labels

loader = DataLoader(dataset, batch_size=8, collate_fn=my_collate)

Common num_workers Guidelines

System           Suggested num_workers
────────────────────────────────────────
Windows (dev)    0  (avoid multiprocessing issues)
Linux / macOS    2–4
High-end server  4–8 (or number of CPU cores / 2)
Google Colab     2

Summary

DataLoader wraps a Dataset and delivers batched, shuffled samples to your training loop. Set shuffle=True for training and False for validation. Use num_workers to load data in parallel for faster pipelines. Set pin_memory=True when training on GPU. Use drop_last=True when batch normalization requires a fixed batch size. Implement a custom collate_fn for variable-length sequences. DataLoader is the bridge between your Dataset and your training loop.