One index at a time

The chapter above this lesson names two objects, a Dataset and a DataLoader, which is the right size for a survey. Build the loader and a third one is sitting on it in plain sight. dl.sampler yields single indices. dl.batch_sampler wraps that sampler and yields lists of indices. dl._index_sampler is whichever of the two the loader will actually iterate, and the choice between them is made once, at construction.

Watch what a Dataset receives and the division of labour is unambiguous. Four indices go out, four separate __getitem__ calls come back, and the stacking happens somewhere else entirely. Your dataset never learns what batch_size is.

That sounds like a detail until you try to make a dataset faster. Caching per batch, reading a contiguous slab off disk once, batching a database query: none of it fits, because the only question your dataset is asked is about one index.

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

class Counted(Dataset):
    def __init__(self, n): self.n, self.asked = n, []
    def __len__(self): return self.n
    def __getitem__(self, i):
        self.asked.append(i)
        return torch.full((3,), float(i))

ds = Counted(10)
dl = DataLoader(ds, batch_size=4)
it = iter(dl)
print("sampler        ", type(dl.sampler).__name__)
print("batch_sampler  ", type(dl.batch_sampler).__name__)
print("_index_sampler ", type(dl._index_sampler).__name__)
print("auto_collation ", dl._auto_collation)
print("collate_fn     ", dl.collate_fn.__name__)
b = next(it)
print("asked for      ", ds.asked)
print("batch shape    ", tuple(b.shape))

# ---- stdout ----
# sampler         SequentialSampler
# batch_sampler   BatchSampler
# _index_sampler  BatchSampler
# auto_collation  True
# collate_fn      default_collate
# asked for       [0, 1, 2, 3]
# batch shape     (4, 3)

Everything the last section showed lives in one small class. _MapDatasetFetcher.fetch takes the list of indices, runs the list comprehension, and hands the result to collate_fn. There is no buffering, no reordering and no threading in it. Read it once and the single-process path stops being mysterious.

The branch above the comprehension is the escape hatch for exactly the problem the last section named. If your dataset defines __getitems__, plural, the fetcher calls it once with the whole index list instead of calling __getitem__ per index. That is the hook a database-backed or memory-mapped dataset wants, and Subset already implements it, which is why wrapping a dataset in Subset does not silently lose the fast path.

The other subclass in the same file handles iterable-style datasets and does something different enough to be worth naming now. It ignores the indices entirely and calls next() on a stored iterator that many times. The third lesson in this arc is about what that costs once workers exist.

class _MapDatasetFetcher(_BaseDatasetFetcher):
    def fetch(self, possibly_batched_index):
        if self.auto_collation:
            if hasattr(self.dataset, "__getitems__") and self.dataset.__getitems__:
                data = self.dataset.__getitems__(possibly_batched_index)
            else:
                data = [self.dataset[idx] for idx in possibly_batched_index]
        else:
            data = self.dataset[possibly_batched_index]
        return self.collate_fn(data)

# a Dataset that also defines __getitems__, batch_size=4, stdout:
# calls: [('getitems', [0, 1, 2, 3])] -> batch [0, 1, 2, 3]
# Subset has __getitems__: True

auto_collation is a property, not an argument, and it is true exactly when the loader has a batch sampler. Pass batch_size=None and it goes false, the batch sampler disappears, and _index_sampler becomes the plain sampler again.

The consequence lands inside fetch, on the else branch. With auto-collation off, the fetcher calls self.dataset[idx] with a single index and passes the result straight to collate_fn, which now defaults to default_convert rather than default_collate. Batching has not been turned off; it has been handed to you. This is the mode an iterable dataset that already yields whole batches wants, and it is also the mode people land in by accident when they set batch_size=None expecting the loader to guess.

One measured detail from the same run is worth carrying into the next section. A Dataset returning (tensor, float) comes back as a float32 input and a float64 label, and nothing in the loader warned about the mismatch.

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

class Pairs(Dataset):
    def __len__(self): return 6
    def __getitem__(self, i): return torch.full((2,), float(i)), float(i)

dl = DataLoader(Pairs(), batch_size=3)
xb, yb = next(iter(dl))
print("label dtype    ", yb.dtype, "| input dtype", xb.dtype)

class Asked(Dataset):
    def __init__(self): self.asked = []
    def __len__(self): return 6
    def __getitem__(self, i):
        self.asked.append(i); return torch.tensor([float(i)])

ds = Asked()
dl0 = DataLoader(ds, batch_size=None)
print("batch_size None: auto_collation", dl0._auto_collation, "| index_sampler", type(dl0._index_sampler).__name__)
b = next(iter(dl0))
print("asked          ", ds.asked, "| got shape", tuple(b.shape))

# ---- stdout ----
# label dtype     torch.float64 | input dtype torch.float32
# batch_size None: auto_collation False | index_sampler SequentialSampler
# asked           [0] | got shape (1,)

default_collate is not a stacking function with special cases bolted on. It is a recursive type walk over a registry, and the registry is a public dictionary you can read at runtime. Look up the type of the first element; if the registry has a handler, call it; otherwise fall through to the structural cases, which recurse.

The structural cases are where the shape of your batch is decided. A mapping is rebuilt as the same mapping type with each key collated across the batch. A namedtuple keeps its class. A plain sequence gets transposed by zip(*batch) before recursing, so a dataset returning (x, y) gives you a two-element result whose halves are batched separately.

One of those cases surprises people who read the type annotation and stopped. A tuple sample does not come back as a tuple. The source returns a list, with a comment saying Backwards compatibility, so type(batch) is list however carefully your __getitem__ built its tuple.

What collate refuses is as informative as what it accepts. Uneven tensor shapes reach torch.stack and die there, with the stack error rather than a collate error, which is why ragged-length data needs a collate_fn of your own. Dictionaries with different key sets die on a KeyError naming the missing key, because the walk indexes every sample by the first sample's keys.

The float64 label from two sections ago is not inference and not a bug. It is a two-line function in the registry, and it says the dtype outright rather than letting torch.tensor guess.

Trace the consequence forward. A Dataset that returns a Python float per sample produces a float64 target, which meets a float32 prediction at the loss. On this machine mse_loss accepted the pair and returned a float64 loss, so the promotion travels silently into the backward pass. The museum's class-targets-as-floats exhibit is the loud version of the same category of mistake; this is the quiet one, and the fix is upstream of both, in what __getitem__ returns.

The registry is also the extension point. default_collate_fn_map is a module-level dict, and the docstring in collate.py shows updating it in place to change how a type batches everywhere. Reach for a custom collate_fn when the batching logic is local; reach for the map when a type should batch the same way across a whole codebase.

def collate_float_fn(batch, *, collate_fn_map: Optional[Dict[Union[Type, Tuple[Type, ...]], Callable]] = None):
    return torch.tensor(batch, dtype=torch.float64)


def collate_int_fn(batch, *, collate_fn_map: Optional[Dict[Union[Type, Tuple[Type, ...]], Callable]] = None):
    return torch.tensor(batch)

default_collate_fn_map: Dict[Union[Type, Tuple[Type, ...]], Callable] = {torch.Tensor: collate_tensor_fn}
with contextlib.suppress(ImportError):
    import numpy as np
    # For both ndarray and memmap (subclass of ndarray)
    default_collate_fn_map[np.ndarray] = collate_numpy_array_fn
    # See scalars hierarchy: https://numpy.org/doc/stable/reference/arrays.scalars.html
    # Skip string scalars
    default_collate_fn_map[(np.bool_, np.number, np.object_)] = collate_numpy_scalar_fn
default_collate_fn_map[float] = collate_float_fn
default_collate_fn_map[int] = collate_int_fn
default_collate_fn_map[str] = collate_str_fn
default_collate_fn_map[bytes] = collate_str_fn