PyTorch Datasets and DataLoaders: Batching, Shuffling and Transforms

Dataset Responsibilities

Models do not train on files directly. They train on batches of tensors. PyTorch separates dataset logic from batching logic: a Dataset knows how to return one sample, and a DataLoader knows how to batch, shuffle, and load many samples efficiently.

A clean data pipeline makes experiments reliable. It should handle train/validation splits, transforms, labels, missing data, and batch shapes predictably.

A Dataset should be simple: store references to examples, implement `__len__`, and implement `__getitem__`. Avoid putting training logic inside the dataset.

Load or reference one sample at a time.
Apply transforms needed for that sample.
Return tensors and labels in a consistent format.

DataLoader Responsibilities

A DataLoader handles batching. It can shuffle samples, use multiple workers, pin memory for GPU transfer, and call a custom collate function when samples have variable lengths.

Use shuffle=True for training, usually False for validation.
Start with num_workers=0 while debugging.
Use custom collate functions for variable-size sequences or nested data.

Custom Tabular Dataset

This example converts NumPy-style arrays into a dataset suitable for regression or classification.

Custom Tabular Dataset

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

class TabularDataset(Dataset):
    def __init__(self, features, labels):
        self.features = torch.tensor(features, dtype=torch.float32)
        self.labels = torch.tensor(labels, dtype=torch.long)

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

    def __getitem__(self, index):
        return self.features[index], self.labels[index]

X = torch.randn(1000, 12).numpy()
y = torch.randint(0, 3, (1000,)).numpy()

dataset = TabularDataset(X, y)
train_ds, val_ds = random_split(dataset, [800, 200])

train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=64, shuffle=False)

features, labels = next(iter(train_loader))
print(features.shape)  # [32, 12]
print(labels.shape)    # [32]

Labels for CrossEntropyLoss should be class indices, not one-hot vectors.
Validation loaders usually do not shuffle because metric order does not matter.

Load the Final Incomplete Batch Deliberately

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

features = torch.arange(10, dtype=torch.float32).reshape(5, 2)
labels = torch.tensor([0, 1, 0, 1, 0])
loader = DataLoader(TensorDataset(features, labels), batch_size=2, shuffle=False)

for batch_features, batch_labels in loader:
    print(tuple(batch_features.shape), batch_labels.tolist())

Output

(2, 2) [0, 1]
(2, 2) [0, 1]
(1, 2) [0]

Before you move on