DataLoader num_workers Bus Error — Docker shm Fix

PyTorch DataLoader and Datasets decouple data storage from batching logic, enabling scalable pipelines that keep GPUs fully utilised. The Dataset class abstracts how individual samples are accessed — one at a time, lazily, from disk or a database. The DataLoader handles batching, shuffling, and multi-process loading on top of whatever Dataset you hand it.

The core problem these tools solve: training on datasets that do not fit in memory while keeping the GPU fed continuously. If data loading is slower than GPU computation, the GPU sits idle between batches — this is called data starvation and it is one of the most common reasons a training run is 3x slower than it should be. The DataLoader solves this by pre-fetching batches in parallel worker processes while the GPU processes the current batch. That overlap is the entire point.

The architectural separation is deliberate and worth internalising early: Dataset knows how to access one sample. DataLoader knows how to batch, shuffle, and parallelise. This means you can swap your data source (disk, SQL, S3, Kafka) without touching the DataLoader, and you can tune the DataLoader's parallelism without touching the Dataset. Each side has one job.

The most common production failure I see in 2026 is the same one I saw in 2022: developers set num_workers=0 during prototyping because it is simpler, everything works, and then they deploy to a real dataset and discover training is 3–5x slower than it needs to be because data loading is serialised on the main thread. The fix is always num_workers >= 1 with pin_memory=True for GPU training — and documenting that requirement so it does not get reverted in a future PR.

What Is PyTorch DataLoader and Datasets and Why Does It Exist?

PyTorch DataLoader and Datasets exist to solve a single concrete problem: how do you train on data that is too large to fit in memory, while keeping a GPU that costs thousands of dollars per hour fully utilised?

The Dataset class — specifically the Map-style variant — requires implementing two methods: __len__ (how many samples exist) and __getitem__ (fetch one sample by index). That is the entire contract. The Dataset knows nothing about batching, shuffling, or parallelism. It just answers 'give me sample 4,217' as fast as it can.

The DataLoader wraps that Dataset and adds everything else: it selects a batch of indices (optionally shuffled), hands those indices to worker processes that call __getitem__ in parallel, collates the results into a batch tensor, and optionally pre-pins that tensor in page-locked memory for faster GPU transfer. The training loop then pulls pre-fetched batches from a queue without waiting.

The performance insight that changes how you think about this: with num_workers=4 and pin_memory=True, the DataLoader is pre-fetching batch N+1 and N+2 while the GPU is still processing batch N. That pipeline overlap is what keeps GPU utilisation above 90%. Without it — with num_workers=0 — every batch is loaded synchronously on the main thread after the GPU finishes the previous one. The GPU sits idle for however long loading takes. On a dataset of real images with augmentations, that idle time can represent 60–70% of wall-clock training time.

As of 2026, with models being trained on increasingly large datasets and GPUs being increasingly expensive, getting this right is not an optimisation — it is table stakes.

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

# A minimal but complete custom Dataset implementation
# This is the pattern you will replicate for every new data source
class ForgeProjectDataset(Dataset):
    def __init__(self, data_list: list, labels: list):
        # Store only metadata in __init__ — never load actual data here
        # If you load data in __init__, it all lands in RAM before training starts
        self.data = data_list
        self.labels = labels

    def __len__(self) -> int:
        # DataLoader uses this to know how many batches constitute one epoch
        return len(self.data)

    def __getitem__(self, idx: int):
        # This is called once per sample, in parallel across num_workers processes
        # Keep it fast: one file read, one transform, return one sample
        sample = torch.tensor(self.data[idx], dtype=torch.float32)
        label  = torch.tensor(self.labels[idx], dtype=torch.long)
        return sample, label


# Minimal working example — four samples, two features each
raw_data   = [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]]
raw_labels = [0, 1, 0, 1]

forge_ds = ForgeProjectDataset(raw_data, raw_labels)

# Production-grade DataLoader configuration for GPU training
forge_loader = DataLoader(
    dataset=forge_ds,
    batch_size=2,
    shuffle=True,           # Reshuffle at every epoch for better generalisation
    num_workers=2,          # Two CPU processes load data in parallel
    pin_memory=True,        # Pre-pin batches in page-locked memory for faster GPU transfer
    persistent_workers=True # Keep workers alive between epochs — avoids 5-10s restart cost
)

# Verify the DataLoader works before starting a long training run
for batch_idx, (samples, labels) in enumerate(forge_loader):
    print(f"Batch {batch_idx}: samples shape {samples.shape}, labels {labels}")
    if batch_idx >= 1:
        break  # Just checking the first two batches

Enterprise Integration: SQL-Backed Datasets

In real production environments, your training data rarely lives in a flat folder of files. It lives in a database — with labels, metadata, train/val/test splits, and versioning all managed in SQL. Implementing a Dataset that queries a SQL backend is one of the more underrated patterns in production ML engineering.

The approach: in __init__, run a single SQL query to fetch metadata only — sample IDs, file paths on disk or object storage, and labels. Store that metadata in memory as a list or DataFrame. In __getitem__, use the file path from metadata to load the actual binary data — the image, audio file, or feature array — from disk or S3. This keeps memory usage proportional to the number of samples (a few bytes per row of metadata), not the size of the data (potentially gigabytes).

The production benefit that makes this pattern worth the setup: when you add new training data, you insert a row into the SQL table and drop the corresponding file on disk. The next training run picks it up automatically via the __init__ query. There is no CSV file to regenerate, no manifest to sync, and no risk of the file list drifting from the actual filesystem state. I have seen teams spend days debugging training regressions that turned out to be a stale CSV pointing to deleted files — this pattern eliminates that entire class of issue.

One thing to watch: do not query SQL inside __getitem__. SQL connections are not thread-safe and cannot be pickled for multi-process workers. Fetch all metadata once in __init__ and do all disk or object-storage I/O in __getitem__.

-- Fetch sample metadata for Dataset __init__
-- We fetch IDs, paths, and labels here — NOT the binary data
-- Binary data is loaded lazily in __getitem__ using the file_path
-- This query runs once at the start of training, not per batch

SELECT
    sample_id,
    file_path,      -- path to the file on NVMe SSD or S3
    label_id,
    split_tag       -- 'train', 'val', or 'test'
FROM io.thecodeforge.training_data
WHERE project_tag = 'vision_v2'
  AND split_tag   = 'train'
  AND is_verified = TRUE   -- exclude samples flagged as corrupted during QA
ORDER BY sample_id ASC;

-- Expected result: one lightweight metadata row per sample
-- Actual image binary data stays on disk until __getitem__ loads it

Containerised Data Pipelines with Docker

Wrapping your training environment in Docker is the standard way to ensure the data pipeline behaves identically across a developer's laptop, a CI server, and a production GPU cluster. It also surfaces the most common PyTorch DataLoader configuration mistake before it costs you a four-hour training run.

The critical Docker configuration that almost everyone gets wrong the first time: when num_workers > 0, PyTorch uses shared memory at /dev/shm to transfer tensors between worker processes and the main process. Docker's default shared memory allocation is 64MB — a sensible default for containerised web services that never heard of PyTorch. For a training job with num_workers=4 and any real batch size, that 64MB fills up within a few epochs and the container dies with a Bus error and no Python traceback. The fix is one flag: --shm-size=2g.

The deployment checklist I use for every new training container: set --shm-size=2g or larger; mount the data directory as a Docker volume rather than copying it into the image (datasets are too large for image layers and change too frequently); pin the PyTorch version explicitly rather than using pytorch/pytorch:latest (latest changes under you in ways that are hard to reproduce); set num_workers based on the CPU cores allocated to the container, not the host machine's total CPU count; and add a pre-flight health check that verifies /dev/shm has enough free space before the training job starts.

The --shm-size flag also belongs in your docker-compose.yml, your Kubernetes pod spec under resources, and your CI job definition — anywhere the container is launched. If it lives in only one place, it will be dropped in a refactor and you will spend an afternoon diagnosing a Bus error that you already fixed six months ago.

# Pin a specific PyTorch release — 'latest' changes under you
# and reproducing a training run six months later becomes impossible
FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime

WORKDIR /app

# Install dependencies before copying source code
# This layer is cached as long as requirements.txt does not change
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

# IMPORTANT: This container MUST be started with --shm-size=2g
# PyTorch DataLoader uses /dev/shm to transfer tensors between worker processes
# Docker's default 64MB causes Bus error crashes when num_workers > 0
# Example: docker run --shm-size=2g --gpus all -v /data:/data thecodeforge/training:latest

# Pre-flight check: verify shared memory is sufficient before training starts
# Exits with a clear error rather than a cryptic Bus error mid-epoch
HEALTHCHECK --interval=10s --timeout=5s --retries=1 \
  CMD python -c "import shutil; free = shutil.disk_usage('/dev/shm').free; assert free > 1e9, f'Insufficient /dev/shm: {free/1e6:.0f}MB free, need 1000MB+'"

CMD ["python", "ForgeDataset.py"]

Common Mistakes and How to Avoid Them

Most DataLoader bugs in production fall into a small set of patterns. Knowing them in advance means you spend time training models instead of debugging pipelines.

The performance mistakes: num_workers=0 is the biggest one — it serialises every sample load on the main thread and GPU sits idle while it happens. Loading data in __init__ instead of __getitem__ is the second — it turns a lazy-loading Dataset into a greedy RAM consumer that OOMs before training even starts.

The correctness mistakes: passing unpickleable objects (open file handles, database connections, lambda functions) to a Dataset when num_workers > 0. Python's multiprocessing pickles the Dataset to send it to each worker process. If any attribute cannot be pickled, the worker hangs silently or crashes without a useful traceback. The fix is to initialise those objects inside __getitem__ (called per sample in each worker) or use worker_init_fn to set them up once per worker process.

The subtle one that costs teams debugging time: forgetting drop_last=True on the training DataLoader. The last batch of an epoch almost always has fewer samples than the configured batch_size. For most loss functions this is harmless, but for BatchNorm it is not — BatchNorm uses batch statistics during training, and a batch of size 1 produces undefined variance. Setting drop_last=True discards the last incomplete batch and ensures consistent batch sizes throughout training. For validation DataLoaders, use drop_last=False — you want to evaluate on every sample, no exceptions.

import torch
from torch.utils.data import DataLoader

# Production-grade DataLoader configuration for GPU training
# Each parameter here solves a specific real problem
loader = DataLoader(
    forge_ds,
    batch_size=32,
    shuffle=True,              # Reshuffle every epoch for better generalisation
    num_workers=4,             # 4 parallel CPU processes — eliminates GPU starvation
    pin_memory=True,           # Pre-pin batches in page-locked memory for faster DMA transfer
    drop_last=True,            # Drop incomplete final batch — prevents BatchNorm issues
    persistent_workers=True,   # Keep workers alive between epochs — avoids 5-10s restart cost
    prefetch_factor=2,         # Each worker pre-fetches 2 batches ahead — reduces wait time
)

# Validation DataLoader — different settings for a reason
val_loader = DataLoader(
    val_ds,
    batch_size=64,             # Larger batch is fine for inference — no gradient storage
    shuffle=False,             # Do NOT shuffle validation — reproducible evaluation
    num_workers=4,
    pin_memory=True,
    drop_last=False,           # Evaluate on EVERY sample — no exceptions
    persistent_workers=True,
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

for samples, labels in loader:
    # non_blocking=True overlaps the CPU-to-GPU transfer with other CPU work
    # Only effective when pin_memory=True is also set
    samples = samples.to(device, non_blocking=True)
    labels  = labels.to(device, non_blocking=True)
    # ... training logic ...

# Random seed consistency across workers
# Without this, all workers inherit the same seed and produce identical augmentations
def worker_init_fn(worker_id: int) -> None:
    import numpy as np
    base_seed = torch.initial_seed() % (2 ** 32)
    np.random.seed(base_seed + worker_id)
    torch.manual_seed(base_seed + worker_id)

# Apply it to any DataLoader that uses random augmentations in __getitem__
auged_loader = DataLoader(
    forge_ds,
    batch_size=32,
    num_workers=4,
    worker_init_fn=worker_init_fn,
    pin_memory=True,
)

Custom collate_fn: Handling Variable-Length Data

The default collate_fn expects every sample in a batch to have the same shape so it can stack them into a uniform tensor. This assumption breaks the moment you work with NLP sequences of different lengths, graphs with different numbers of nodes, or images that have not been resized to a fixed resolution.

A custom collate_fn lets you define exactly how a list of heterogeneous samples becomes a batch. The most common pattern — one you will write or encounter in almost every NLP project — is pad-and-mask: pad all sequences to the length of the longest sequence in the batch, and return a binary mask tensor that tells downstream layers which positions are real data and which are padding. Attention layers, loss functions, and pooling operations all need this mask to avoid treating padding as signal.

The production subtlety that trips people up: collate_fn runs on the main thread, not inside the worker processes. This means even with num_workers=4, a slow collate_fn becomes the bottleneck for the entire pipeline. Keep it to reshaping and padding only. If you find yourself sorting sequences, computing complex statistics, or doing any non-trivial transformation in collate_fn, move that work into __getitem__ where it can run in parallel across workers.

For NLP work in 2026, most teams use Hugging Face's DataCollatorWithPadding which implements this pattern with tokeniser-aware padding. But understanding the underlying collate_fn contract means you can customise it when the standard collators do not fit your data structure.

import torch
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
from typing import List, Tuple


def collate_variable_length(
    batch: List[Tuple[torch.Tensor, torch.Tensor]]
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Pads variable-length sequences to the longest in the batch.
    Returns:
        padded   — (batch_size, max_len, feature_dim) padded sequences
        labels   — (batch_size,) classification labels
        mask     — (batch_size, max_len) float mask: 1.0 real, 0.0 padding

    Runs on the main thread — keep this function fast.
    Heavy transforms belong in __getitem__, not here.
    """
    sequences, labels = zip(*batch)

    # pad_sequence pads to the longest sequence in this batch
    # batch_first=True gives shape (batch, seq_len, features)
    padded = pad_sequence(sequences, batch_first=True, padding_value=0.0)

    # Build the mask: 1.0 where real data, 0.0 where padding
    # Downstream attention layers and loss functions use this to ignore padding
    lengths = torch.tensor([len(s) for s in sequences], dtype=torch.long)
    mask    = torch.zeros(padded.shape[0], padded.shape[1], dtype=torch.float32)
    for i, length in enumerate(lengths):
        mask[i, :length] = 1.0

    labels_stacked = torch.stack(labels)
    return padded, labels_stacked, mask


# Toy variable-length dataset
class VariableLengthDataset(Dataset):
    def __init__(self, num_samples: int = 200, feature_dim: int = 128):
        self.num_samples = num_samples
        self.feature_dim = feature_dim

    def __len__(self) -> int:
        return self.num_samples

    def __getitem__(self, idx: int):
        # Sequence length varies per sample — this is what breaks the default collate_fn
        seq_len  = torch.randint(10, 60, (1,)).item()
        features = torch.randn(seq_len, self.feature_dim)
        label    = torch.tensor(idx % 2, dtype=torch.long)  # binary label
        return features, label


variable_length_dataset = VariableLengthDataset(num_samples=200, feature_dim=128)

loader = DataLoader(
    variable_length_dataset,
    batch_size=32,
    shuffle=True,
    collate_fn=collate_variable_length,  # custom collation for variable-length sequences
    num_workers=4,
    pin_memory=True,
)

for batch_idx, (padded_seqs, labels, mask) in enumerate(loader):
    # padded_seqs: (32, max_len_in_batch, 128)
    # mask:        (32, max_len_in_batch) — 1.0 for real tokens, 0.0 for padding
    real_token_count = mask.sum().item()
    total_positions  = mask.numel()
    padding_pct      = 100 * (1 - real_token_count / total_positions)
    print(f"Batch {batch_idx}: shape {padded_seqs.shape} | "
          f"padding {padding_pct:.1f}% | labels {labels[:4]}")
    if batch_idx >= 2:
        break

IterableDataset: Streaming Data Without Random Access

Map-style datasets require __len__ and __getitem__ — random access to any sample by index. This breaks when data arrives as a stream (Kafka, network logs, real-time sensor feeds) or when the dataset is genuinely too large to index. IterableDataset solves this by yielding samples sequentially without needing to know the total size.

The use case that justifies reaching for IterableDataset: training on a live event stream where the concept of 'total dataset size' does not exist, or on a dataset so large that generating a complete index would take longer than training itself. The DataLoader iterates through the __iter__ method, batches samples as they arrive, and provides limited shuffling within a buffer of recent samples.

The production trade-off that you need to understand before choosing IterableDataset: it cannot shuffle globally because it never knows the full dataset. It can shuffle within a configurable buffer of recent samples, but the model always sees data in approximately the order it arrives in the stream. If the stream has any temporal structure — and real data almost always does — the model will see a biased distribution. For training data that can be indexed, Map-style datasets with global shuffling are strictly better. Use IterableDataset only when indexing is genuinely impossible.

One non-obvious operational issue with num_workers > 0 and IterableDataset: each worker receives its own copy of the __iter__ method and will iterate the entire stream independently. Without sharding the stream across workers, every sample gets loaded num_workers times. You need to detect the worker ID inside __iter__ using torch.utils.data.get_worker_info() and partition the stream so each worker handles a distinct subset.

import torch
from torch.utils.data import IterableDataset, DataLoader, get_worker_info
from typing import Iterator, Tuple


class StreamingSensorDataset(IterableDataset):
    """
    Yields samples from a simulated sensor stream.
    In production, replace __iter__ with a Kafka consumer,
    network socket reader, or database cursor.

    IMPORTANT: with num_workers > 0, each worker calls __iter__ independently.
    Without sharding, every sample is loaded num_workers times.
    The __iter__ below handles worker partitioning automatically.
    """

    def __init__(self, num_samples: int = 10000):
        self.num_samples = num_samples

    def __iter__(self) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
        worker_info = get_worker_info()

        if worker_info is None:
            # Single-process loading — yield the full stream
            start, end = 0, self.num_samples
        else:
            # Multi-process loading — partition stream across workers
            # Each worker gets a non-overlapping slice of the sample range
            per_worker = self.num_samples // worker_info.num_workers
            worker_id  = worker_info.id
            start      = worker_id * per_worker
            end        = start + per_worker if worker_id < worker_info.num_workers - 1 else self.num_samples

        for i in range(start, end):
            # Simulate streaming sensor data — replace with real I/O in production
            features = torch.randn(10)
            label    = torch.tensor(float(features[0] > 0), dtype=torch.float32)
            yield features, label


stream_ds     = StreamingSensorDataset(num_samples=10000)
stream_loader = DataLoader(
    stream_ds,
    batch_size=64,
    num_workers=2,  # Each worker now handles a non-overlapping shard of the stream
)

for batch_idx, (features, labels) in enumerate(stream_loader):
    if batch_idx >= 5:
        break
    print(f"Batch {batch_idx}: features {features.shape}, "
          f"label distribution: {labels.sum().item():.0f}/{len(labels)} positive")

Why Batch Size Breaks Your Model (And How to Pick It)

Most tutorials tell you batch size is a hyperparameter you tune for memory. That's half the story. The real reason batch size matters is gradient variance.

Small batches (32-64) give noisy gradients. That noise acts as regularization — it helps escape sharp minima and generalizes better. Large batches (512+) average out the noise, giving you cleaner gradients that converge faster but often to flatter minima with worse generalization. This isn't theory. It's why you see models train slower per epoch with small batches but reach better test accuracy.

Here's the trap: doubling batch size to speed up training doesn't halve wall time. DataLoader overhead, GPU memory bandwidth, and the cost of synchronizing gradients across devices grow non-linearly. A batch of 1024 isn't 32x faster than a batch of 32. It's maybe 2-3x faster per epoch, but you'll need more epochs to converge.

Rule of thumb: start at 32 for vision tasks, 16 for NLP. Double only if gradients look stable and validation loss isn't plateauing early. Monitor GPU utilization with nvidia-smi — if you're below 80%, increase batch size until you hit 90%.

// io.thecodeforge — ml-ai tutorial

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

def profile_batch_sizes(model, inputs, labels, batch_sizes=[32, 64, 128, 256]):
    dataset = TensorDataset(inputs, labels)
    results = {}
    for batch_size in batch_sizes:
        loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4)
        start = time.perf_counter()
        for x, y in loader:
            out = model(x)
            loss = out.sum()
            loss.backward()
        elapsed = time.perf_counter() - start
        results[batch_size] = round(elapsed, 2)
    return results

# dummy data: 10000 samples, 256 features
dummy_inputs = torch.randn(10000, 256)
dummy_labels = torch.randint(0, 10, (10000,))
model = torch.nn.Linear(256, 10)

profiled = profile_batch_sizes(model, dummy_inputs, dummy_labels)
for bs, t in profiled.items():
    print(f"batch_size={bs}: {t}s per epoch")

Shuffling Isn't Optional — Here's Why Your Model Memorizes Without It

You've seen the shuffle=True argument in DataLoader. Ever wondered what happens when you set it to False? Your model learns the order of your data, not the patterns. This is called dataset memorization, and it's insidious.

Consider a dataset sorted by class: first 1000 cats, then 1000 dogs. Without shuffling, the first epoch shows the model only cats. The model adapts to predict 'cat' for everything. Then epoch 2 shows only dogs. The model flips to predict 'dog'. This creates oscillations in loss that never converge. Even if your data isn't explicitly sorted, inherent ordering from data collection (e.g., timestamps, sensor IDs) leaks bias.

Shuffling breaks temporal correlations between samples. It ensures each mini-batch is an i.i.d. sample from the data distribution. This stabilizes gradient updates and prevents the model from exploiting spurious order patterns.

But there's a catch: shuffling entire datasets is expensive. For datasets larger than RAM (think 100GB+), full shuffling is impossible. That's when you use Sampler objects — specifically RandomSampler with replacement or distributed-aware samplers for multi-GPU training. Never implement your own shuffle logic. PyTorch's DataLoader does it correctly with Fisher-Yates in C++. Your Python loop will be 100x slower.

// io.thecodeforge — ml-ai tutorial

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

# Simulate sorted data: class 0 then class 0 then class 1
data = torch.cat([torch.zeros(1000, 10), torch.ones(1000, 10)], dim=0)
labels = torch.cat([torch.zeros(1000), torch.ones(1000)], dim=0).long()
dataset = TensorDataset(data, labels)

loader_no_shuffle = DataLoader(dataset, batch_size=64, shuffle=False)
loader_shuffle = DataLoader(dataset, batch_size=64, shuffle=True)

for name, loader in [("No shuffle", loader_no_shuffle), ("Shuffle", loader_shuffle)]:
    class_counts = torch.zeros(2)
    for _, y in loader:
        for c in range(2):
            class_counts[c] += (y == c).sum().item()
    print(f"{name}: batch class distribution: {class_counts.tolist()}")

When DataLoader Bottlenecks Your GPU: Num Workers and Prefetch

Your GPU sits at 30% utilization while training. You blame the model. I blame your DataLoader. The culprit is almost always insufficient num_workers or missing prefetch_factor. Here's the math.

By default, DataLoader uses num_workers=0 — meaning the main process loads data serially. That's a death sentence for GPU-bound workloads. Each time the GPU finishes a batch, it waits for the CPU to load the next one. This creates a bubble of idle GPU time.

Set num_workers to the number of CPU cores (not threads). For a 16-core machine, use 12-14 workers. Each worker prefetches batches independently. The prefetch_factor (default 2) controls how many batches each worker queues ahead. Increase it to 4 or 8 if your data transforms are heavy (e.g., image augmentations).

But there's a ceiling. Too many workers cause contention on disk I/O and memory bandwidth. Monitor with htop and nvidia-smi. If CPU usage hits 100% or disk reads saturate, back off workers. Also, Windows users: num_workers=0 is your only safe option due to multiprocessing quirks — use Linux for serious training.

One more thing: when using custom transforms, move heavy operations (resizing, normalization) into the Dataset's __getitem__ and let workers parallelize them. Don't apply them after batching — that's serial and slow.

// io.thecodeforge — ml-ai tutorial

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

# Simulate heavy data transform: 100ms per sample
def heavy_load(x):
    time.sleep(0.1)
    return x * 2

data = torch.randn(500, 256)
labels = torch.randint(0, 10, (500,))
dataset = TensorDataset(data, labels)

for workers in [0, 2, 4, 8]:
    loader = DataLoader(
        dataset, batch_size=32, shuffle=True,
        num_workers=workers, prefetch_factor=2
    )
    start = time.perf_counter()
    for _ in range(10):  # 10 epochs
        for x, y in loader:
            x = heavy_load(x)
    elapsed = time.perf_counter() - start
    print(f"num_workers={workers}: {elapsed:.2f}s")

Stop Guessing Transforms: Compose Like a Pro

You don't hand-roll random crops and flips. torchvision.transforms.Compose chains them in one pipeline. The WHY is simple: your model sees invariant features faster when it can't rely on pixel-location memorization. Compose lets you stack RandomHorizontalFlip, ColorJitter, and RandomResizedCrop into a single callable that fits into Dataset.__getitem__.

Here's how it works: you pass a list of transforms to Compose, it applies each sequentially at index time. No manual loop. No duplicated logic. For production, pin the transforms order: normalization must be last, after all random augmentations. If you flip after normalize, your color distribution breaks.

// io.thecodeforge — ml-ai tutorial

import torch
from torchvision import transforms
from torchvision.datasets import ImageFolder

# Compose defines the pipeline order
pipeline = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

dataset = ImageFolder(root='./data/train', transform=pipeline)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32)

for images, labels in dataloader:
    print(images.shape)  # torch.Size([32, 3, 224, 224])
    print(labels.shape)  # torch.Size([32])
    break

Don't Import ImageFolder Blind — Know When to Dump It

Waste of time: asking "what's the difference between DataLoader and ImageFolder." They're not even the same thing. ImageFolder is a dataset class that assumes your images sit in subfolders named by class. DataLoader is a load balancer — it wraps any dataset, batches samples, and spins up workers. ImageFolder gives you (PIL.Image, label) tuples. DataLoader gives you tensors.

Here's the real difference: use ImageFolder when your data is already sorted into class folders on disk. That's it. If your data is in a CSV, a SQL table, or a blob store, skip ImageFolder. Write a custom Dataset. ImageFolder is a convenience wrapper over datasets.DatasetFolder. It reads file paths, applies your transform, and returns samples. DataLoader then handles concurrency, batching, and shuffling. Know the boundary.

// io.thecodeforge — ml-ai tutorial

from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

# ImageFolder: reads class subdirectories
dataset = ImageFolder(
    root='./data/train',
    transform=transforms.ToTensor()
)

# DataLoader: handles batch delivery, not data structure
dataloader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4
)

sample_img, label = dataset[0]
print(f"Class index: {label}, Class name: {dataset.classes[label]}")
print(f"Number of classes: {len(dataset.classes)}")

Prefetching Is Free Speed — Use It or Waste GPU Cycles

Your GPU sits idle while DataLoader fetches the next batch. That's the problem. prefetch_factor in DataLoader tells how many batches to load ahead while the current one trains. Default is 2. Bump it to 4 or 8 on high-throughput pipelines. The WHY is simple: overlap I/O with compute. While your model runs forward and backward, workers load the next prefetch_factor * batch_size samples into a queue.

Combine this with num_workers. Rule of thumb: set workers to 4–8 per GPU, then prefetch to 4. Watch your RAM. If you hit swap, back off workers or lower prefetch. Production setup: log time_to_first_batch and iter_time — if iter_time < load_time, increase prefetch. Your GPU will thank you.

// io.thecodeforge — ml-ai tutorial

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

# Synthetic data — real pipeline mirrors this
data = torch.randn(10000, 3, 224, 224)
labels = torch.randint(0, 10, (10000,))
dataset = TensorDataset(data, labels)

# Prefetch 8 batches ahead on each worker
dataloader = DataLoader(
    dataset,
    batch_size=128,
    shuffle=True,
    num_workers=4,
    prefetch_factor=8,
    pin_memory=True  # moves tensors to pinned memory for async GPU transfer
)

start_time = time.time()
for epoch in range(1):
    for batch_idx, (inputs, targets) in enumerate(dataloader):
        if batch_idx == 0:
            print(f"First batch delivered in {time.time()-start_time:.2f}s")
        # Simulate training
        _ = inputs.mean()

Creating an Instance of the Dataset

Most engineers copy-paste dataset instantiation without understanding the memory cost. Every time you create a Dataset object, you are building a mapping from indices to data points. For map-style datasets, this means the entire list of file paths or SQL row pointers is loaded into memory. A common mistake is creating a separate Dataset instance per epoch instead of reusing one — this doubles memory overhead for zero benefit. Instead, instantiate the Dataset once and feed it into DataLoader repeatedly. When using transforms, the Dataset instance stores them as attributes; passing transforms during creation lets you swap augmentations without recreating the entire mapping. If your data is on remote storage, lazy-loading strategies inside the Dataset constructor can prevent OOM errors. Always profile memory usage with torch.cuda.memory_summary() after instantiation to catch leaks early.

// io.thecodeforge — ml-ai tutorial

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import os

class ImageFileDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.files = sorted(os.listdir(root_dir))  # single init
        self.transform = transform

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        img = torch.randn(3, 224, 224)  # placeholder load
        if self.transform:
            img = self.transform(img)
        return img

# ONE instance per run — reuse across epochs
dataset = ImageFileDataset('/data/images', transform=transforms.ToTensor())
dloader = DataLoader(dataset, batch_size=32, shuffle=True)

for epoch in range(10):
    for batch in dloader:  # same dataset, no re-init
        pass

Benefits of Using Mini-Batches

Training on mini-batches isn't just about fitting data into GPU memory — it fundamentally changes optimization dynamics. Single-sample (stochastic) updates have high variance because gradients from one sample don't represent the full data distribution. Full-batch gradients are stable but computationally prohibitive and generalize poorly due to sharp minima. Mini-batches hit the sweet spot: they average gradients over a small, representative subset, reducing variance while maintaining per-step efficiency. PyTorch's DataLoader handles batching automatically via the batch_size parameter, but the real win is parallelization — mini-batches let you exploit matrix operations on GPU hardware, achieving 10–100x throughput over sequential single-sample updates. The optimal batch size depends on your model's memory footprint and the dataset's intrinsic variance; start with powers of 2 between 32 and 512. Large batches require higher learning rates to maintain convergence speed.

// io.thecodeforge — ml-ai tutorial

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

# Simulate dataset
X = torch.randn(10000, 784)
y = torch.randint(0, 10, (10000,))
dataset = TensorDataset(X, y)

# Mini-batch vs full-batch comparison
mini_loader = DataLoader(dataset, batch_size=64, shuffle=True)
full_loader  = DataLoader(dataset, batch_size=10000, shuffle=False)

# Mini-batch: 157 gradient steps per epoch
for batch_x, batch_y in mini_loader:
    loss = (batch_x.sum() - batch_y.sum()).pow(2)
    loss.backward()  # efficient matrix ops
    break

# Full-batch: 1 step per epoch — memory spike
for batch_x, batch_y in full_loader:
    loss = (batch_x.sum() - batch_y.sum()).pow(2)
    loss.backward()  # OOM risk on large GPU
    break