# 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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
inputs, targets = next(iter(train_loader))

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

Typical shapes are:

```text
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.

```python
loader = DataLoader(dataset, batch_size=128)
```

Inspect the number of batches directly with:

```python
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.

```text
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.

```python
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`.

```python
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.

```python
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].

```python
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.

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

```python
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.

```python
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.

```python
weights = ...

transform = weights.transforms()
```

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