PyTorch / Tensor, gradient, and module foundations / Lesson 5

Datasets, loaders, and collation

Define one sample and one batch explicitly while preserving identity, coverage, and split boundaries.

About 4 min readContent reviewed 2026-10-28

On this page

Lesson goals

Implement a small map-style Dataset contract.
Explain DataLoader batching, collation, and final-batch policy.
Keep worker, sampling, and preprocessing assumptions visible.

A dataset describes access to observations

A map-style Dataset commonly supplies __len__ and __getitem__, with each index returning one defined observation. The index is an access mechanism, not necessarily the permanent observation identity. An IterableDataset instead supplies an iteration stream and has different sampling and worker responsibilities. Choose the interface from the data source rather than imitating an unrelated example. This course uses three supplied synthetic rows entirely in memory. It downloads no dataset, reads no private records, and starts no external service. Define feature order, target meaning, and the split manifest outside the loader itself.

A sample contract determines what batching can do

The illustrative Dataset returns a mapping containing a string id, a two-feature float64 CPU tensor, and a scalar target equal to the supplied feature sum. DataLoader’s documented default collation handles mappings recursively, stacks compatible tensors, and preserves non-tensor information according to its rules. The batch x shape becomes (B,2) and y becomes (B,). Variable-length tensors cannot simply be stacked into a rectangular batch without another policy. A custom collate_fn can pad, aggregate, or preserve a list, but must state lengths, masks, and target alignment rather than silently discard information.

python Illustrative CPU PyTorch 2.14 API program; torch unavailable, expected output derived

import torch
from torch.utils.data import Dataset, DataLoader
class TinyRows(Dataset):
    def __init__(self):
        self.rows = [("a",[1.0,2.0]),("b",[3.0,4.0]),("c",[5.0,6.0])]
    def __len__(self):
        return len(self.rows)
    def __getitem__(self, index):
        key, values = self.rows[index]
        return {"id":key, "x":torch.tensor(values,dtype=torch.float64,device="cpu"),
                "y":torch.tensor(sum(values),dtype=torch.float64,device="cpu")}
loader = DataLoader(TinyRows(),batch_size=2,shuffle=False,num_workers=0,drop_last=False)
for batch in loader:
    print(batch["id"],tuple(batch["x"].shape),batch["y"].tolist())

Expected output

['a', 'b'] (2, 2) [3.0, 7.0]
['c'] (1, 2) [11.0]

Default collation stacks equal-shaped tensors and retains string ids in batches; the final batch is smaller because drop_last=False. No DataLoader was executed.

Batching choices change coverage and weighting

The loader fixes shuffle=False, num_workers=0, batch_size=2, and drop_last=False. It therefore describes a two-row batch followed by a one-row batch. Dropping the last batch can exclude observations and change the task’s effective training population. A metric averaged over batch means gives each batch equal weight, which differs from observation weighting when batch sizes vary. The companion verifies this arithmetic on plain lists. Keep counts or numerator-and-denominator totals for the intended metric, and test a dataset length that is not a multiple of the batch size.

python Executed deterministic standard-library companion

ids = ["a","b","c"]
batches = [ids[start:start+2] for start in range(0,len(ids),2)]
recovered = [key for batch in batches for key in batch]
print("batches", batches)
print("coverage preserved", recovered == ids)
means, counts = [2.0,5.0], [2,1]
weighted = sum(mean*count for mean,count in zip(means,counts))/sum(counts)
print(f"sample-weighted mean={weighted:.1f}")

Expected output

batches [['a', 'b'], ['c']]
coverage preserved True
sample-weighted mean=3.0

The executed companion checks deterministic partition coverage and weighting of an incomplete final batch, not tensor collation, sampling, or worker behavior.

Workers and randomness introduce additional state

Using workers can change process boundaries, ordering assumptions, random-state handling, and resource use. Iterable data can be duplicated across workers unless partitioning is designed correctly. The example intentionally uses zero workers and no shuffle to avoid those execution concerns; it does not test multiprocessing. Shuffling changes sample order but does not create a valid train/evaluation split. Fit imputation, scaling, or vocabulary state only within the permitted training boundary, then apply that fitted state in each sample path. A transform inside __getitem__ can still leak if its learned parameters were computed globally beforehand.

Inspect a batch before adding training

Check ids, tensor shapes, dtypes, target alignment, valid lengths, and the handling of missing or unsupported samples. Preserve identity through filtering and any custom collation so predictions can later be joined correctly. Reordering completion must not become reordering labels independently. Include empty, singleton, uneven, and variable-length cases under explicit policies. The framework program remains an unexecuted documented expectation because torch is unavailable. The standard-library companion checks only list partitioning and weighted arithmetic; it does not validate Dataset methods, default_collate, samplers, workers, pinned memory, or data-loader performance.

Common errors and fixes

The last examples vanish from each epoch

Why it happens: drop_last=True or a custom batching loop discarded the final partial batch.
How to fix it: Choose and document the intended coverage policy.
Prevent it next time: Test uneven lengths and reconcile observation ids.

A metric changes when only batch size changes

Why it happens: Batch means were averaged without observation counts.
How to fix it: Aggregate the intended numerator and denominator across batches.
Prevent it next time: Include one full and one partial batch in a known arithmetic fixture.

Workers repeat iterable samples

Why it happens: Each worker traversed the same stream without partitioning.
How to fix it: Use an appropriate worker-aware stream policy in the supported runtime.
Prevent it next time: Test identity coverage and multiplicity rather than only total throughput.

Check your understanding

1. Why does the last x batch have shape (1,2) instead of (2,2)?

Show worked solution

There are three observations and batches hold at most two. With drop_last=False, the remaining observation is retained as a batch of one with two features. It should not lose the batch dimension or acquire a copied duplicate merely to fill a rectangle.

2. Batch means two and five come from counts two and one. What is the observation mean?