Data Management

Data Management#

Use this page to quickly reference how to work with PyTorch datasets, create batches with data loaders, split data reproducibly, and apply preprocessing without leaking information across partitions.

Datasets#

Dataset#

A PyTorch Dataset provides access to individual examples.

  • __len__ returns the number of examples.

  • __getitem__ returns one example at the given index.

from torch.utils.data import Dataset

class ExampleDataset(Dataset):

    def __len__(self):
        return ...

    def __getitem__(self, index):
        inputs = ...
        target = ...
        return inputs, target

The dataset does not group examples into batches. That is the responsibility of the DataLoader.

TensorDataset#

Use TensorDataset when the complete dataset is already stored in tensors.

from torch.utils.data import TensorDataset

dataset = TensorDataset(features, targets)

All tensors passed to TensorDataset must contain the same number of examples along dimension 0.

Inspecting a Dataset#

Inspect examples before training to verify shapes, dtypes, and labels.

print("Number of examples:", len(dataset))

inputs, target = dataset[0]

print("Input shape:", inputs.shape)
print("Input dtype:", inputs.dtype)
print("Target:", target)

For image data, visual inspection is also useful for detecting corrupted inputs or incorrect labels.

Data Loaders#

Batching Terminology#

Term

Meaning

Sample

One example from a dataset

Batch

A collection of examples processed together

Iteration / training step

One parameter update using one batch

Epoch

One complete pass through the training dataset

DataLoader#

A DataLoader retrieves examples from a dataset and groups them into batches.

from torch.utils.data import DataLoader

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

valid_loader = DataLoader(valid_dataset, batch_size=256, shuffle=False)

Training examples are normally shuffled. Validation and test examples do not need shuffling.

Argument

Meaning

dataset

Dataset providing individual examples

batch_size

Number of examples grouped into one batch

shuffle

Randomly reorder examples before batching

drop_last

Discard the final incomplete batch when True

Inspecting a Batch#

Retrieve one batch directly from the loader.

inputs, targets = next(iter(train_loader))

print(inputs.shape)
print(targets.shape)

Typical shapes are:

tabular inputs   → (B, D)
image inputs     → (B, C, H, W)
targets          → (B,)

where B is the current batch size.

Inspecting one batch is one of the fastest ways to detect mismatches between the data pipeline, model input, and loss function.

Batch Size#

Batch size controls how many examples are processed together in one iteration.

loader = DataLoader(dataset, batch_size=128)

Inspect the number of batches directly with:

len(loader)

With drop_last=False, a dataset containing \(N\) examples and batch size \(B\) produces \(\left\lceil \frac{N}{B} \right\rceil\) batches.

Incomplete Final Batches#

If the dataset size is not divisible by the batch size, the final batch is smaller.

1000 examples
batch size 64

→ 15 full batches of 64
→ 1 final batch of 40

This matters when aggregating losses and metrics across an epoch.

Warning

Do not assume that every batch has the same number of examples unless drop_last=True.

Dataset Splits#

random_split#

For a simple random partition, use random_split with a seeded generator.

from torch.utils.data import random_split

total   = len(dataset)
n_train = int(total * 0.8)
lengths = [n_train, total - n_train]

generator = torch.Generator().manual_seed(42)

train_dataset, valid_dataset = random_split(dataset, lengths, generator=generator)

The seed fixes the random partition so that repeated runs assign the same examples to each subset.

Warning

Subsets created with random_split refer to the same underlying dataset. If that dataset owns a transform, all subsets use that transform.

Subset#

When more control is needed, create the split indices explicitly and use Subset.

from torch.utils.data import Subset

total = len(dataset)
n_train = int(total * 0.8)

generator = torch.Generator().manual_seed(42)
indices = torch.randperm(total, generator=generator)

train_indices = indices[:n_train]
valid_indices = indices[n_train:]

train_dataset = Subset(dataset, train_indices)
valid_dataset = Subset(dataset, valid_indices)

Explicit indices are useful when the same partition must be reused across several dataset instances or preprocessing configurations. For example, training and validation images may require different transforms. Create two dataset instances that refer to the same examples, then apply the same split indices to each one.

train_data = ImageFolder(root="data", transform=train_transform)
valid_data = ImageFolder(root="data", transform=eval_transform)

train_dataset = Subset(train_data, train_indices)
valid_dataset = Subset(valid_data, valid_indices)

Structured Splits#

Random example-level splitting is inappropriate when the observations are not independent. Split according to the unit that must remain separate across partitions.

Data structure

Split by

Independent examples

Example

Repeated observations from one subject

Subject or group

Several samples derived from one source

Original source

Time-ordered observations

Time

The split must respect the structure that defines independent examples. Repeated observations from the same subject, derived samples from the same source, or time-ordered observations should not be randomly distributed across training and evaluation partitions.

Preprocessing#

Per-Example Transformations#

Per-example transformations operate independently on each example and do not estimate parameters from the dataset. For example, an image can be converted to a tensor and its values scaled to the range [0, 1].

import torchvision.transforms.v2 as v2

transform = v2.Compose([
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True)
])

Such transformations can be attached directly to many torchvision datasets.

Fitted Preprocessing#

Fitted preprocessing estimates parameters or rules from data. It must be fitted using the training set only and then applied unchanged to validation and test data. Typical examples include:

  • means and standard deviations;

  • imputation values;

  • feature-selection rules;

  • dimensionality-reduction parameters.

Create the data partitions before fitting any data-dependent preprocessing.

Warning

Validation and test data must not influence fitted preprocessing parameters.

Image Data#

ImageFolder#

Use ImageFolder to load images that are organized in one directory per class.

data/
│
├── cat/
│   ├── cat001.jpg
│   ├── cat002.jpg
│   └── ...
│
├── dog/
│   ├── dog001.jpg
│   ├── dog002.jpg
│   └── ...
│
└── bird/
    ├── bird001.jpg
    ├── bird002.jpg
    └── ...

from torchvision.datasets import ImageFolder

dataset = ImageFolder(root="data", transform=transform)

Useful attributes include .classes and .class_to_idx.

Attribute

Meaning

dataset.classes

List of class names in alphabetical order

dataset.class_to_idx

Dictionary mapping class names to integer indices

ImageFolder assigns integer class indices from the directory names.

Training and Evaluation Transforms#

Training and evaluation data may use different transformations.

train_transform = v2.Compose([
    v2.ToImage(),
    v2.RandomHorizontalFlip(),
    v2.ToDtype(torch.float32, scale=True)
])

eval_transform = v2.Compose([
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True)
])

Training may use stochastic transformations such as augmentation. Validation and test transformations should normally be deterministic.

Pretrained Model Transforms#

Pretrained torchvision weights provide the preprocessing expected by the model.

weights = ...

transform = weights.transforms()

Use the transformation associated with the selected pretrained weights unless the application requires a deliberate alternative. See Transfer Learning for freezing, replacing heads, and fine-tuning pretrained models.