diff --git a/lightly/data/__init__.py b/lightly/data/__init__.py index 257626865..ea087b65c 100644 --- a/lightly/data/__init__.py +++ b/lightly/data/__init__.py @@ -25,3 +25,8 @@ ) from lightly.data.dataset import LightlyDataset from lightly.data.ijepa_collate import IJEPAMaskCollator + +# collate and legacy_collate are not re-exported here: the name is taken by the +# deprecated lightly.data.collate module. Import them from lightly.data.sample +# until that module goes. +from lightly.data.sample import Sample, View diff --git a/lightly/data/collate.py b/lightly/data/collate.py index deeefda66..00ce00d75 100644 --- a/lightly/data/collate.py +++ b/lightly/data/collate.py @@ -11,9 +11,14 @@ import torchvision from PIL import Image -from lightly.transforms import GaussianBlur, Jigsaw, RandomSolarization +# Imported from their own modules rather than from the lightly.transforms +# package: lightly.transforms now reaches lightly.data.sample for View, and going +# through the package __init__ closes that loop. +from lightly.transforms.gaussian_blur import GaussianBlur +from lightly.transforms.jigsaw import Jigsaw from lightly.transforms.random_crop_and_flip_with_grid import RandomResizedCropAndFlip from lightly.transforms.rotation import random_rotation_transform +from lightly.transforms.solarize import RandomSolarization from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T from lightly.transforms.utils import IMAGENET_NORMALIZE diff --git a/lightly/data/multi_view_collate.py b/lightly/data/multi_view_collate.py index e89074c83..c718aa111 100644 --- a/lightly/data/multi_view_collate.py +++ b/lightly/data/multi_view_collate.py @@ -12,8 +12,12 @@ class MultiViewCollate: multiple views of an image, a label, and a filename. It outputs these as separate grouped tensors for easy batch processing. + Transforms that have moved to the view contract return views rather than + bare tensors, and ``lightly.data.sample.legacy_collate`` is the equivalent + of this class for them. + Example: - >>> transform = SimCLRTransform() + >>> transform = SimSiamTransform() >>> dataset = LightlyDataset(input_dir, transform=transform) >>> dataloader = DataLoader( ... dataset, batch_size=4, collate_fn=MultiViewCollate() diff --git a/lightly/data/sample.py b/lightly/data/sample.py new file mode 100644 index 000000000..33fed22e0 --- /dev/null +++ b/lightly/data/sample.py @@ -0,0 +1,202 @@ +"""The batch contract: a sample is a list of typed views.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Sequence + +import torch +from torch import Tensor + +__all__ = ["Sample", "View", "collate", "legacy_collate"] + + +@dataclass +class View: + """One view of one sample, labelled with what it is. + + A single item holds ``data`` at ``(C, H, W)``. After collation the same field + holds ``(B, C, H, W)``. + + Attributes: + data: + The view itself. + stream: + The modality the view came from: ``image``, ``text``, ``audio``, + ``state`` or ``action``. + role: + What the method does with the view: ``view``, ``global``, ``local``, + ``context``, ``target`` or ``anchor``. + extras: + Whatever the transform emitted alongside the view, such as a mask, a + grid or patch ids. Collation stacks every entry. + """ + + data: Tensor + stream: str = "image" + role: str = "view" + extras: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class Sample: + """A batch of views, plus what belongs to the sample rather than to a view. + + Attributes: + views: + The views, in the order the transform produced them. + meta: + Per-sample values such as the target, the filename or an episode id. + """ + + views: list[View] + meta: dict[str, Any] = field(default_factory=dict) + + def by_role(self, role: str) -> list[View]: + """Returns the views with the given role, in view order. + + Args: + role: The role to select. + + Returns: + The matching views, empty if there are none. + """ + return [view for view in self.views if view.role == role] + + def by_stream(self, stream: str) -> list[View]: + """Returns the views with the given stream, in view order. + + Args: + stream: The stream to select. + + Returns: + The matching views, empty if there are none. + """ + return [view for view in self.views if view.stream == stream] + + +def _stack(values: Sequence[Any]) -> Any: + """Stacks tensors, leaves anything else as a list.""" + if all(isinstance(value, Tensor) for value in values): + return torch.stack(list(values)) + return list(values) + + +def _merge(views: Sequence[View], position: int) -> View: + """Merges the view at one position across the samples of a batch. + + The collate only sees sample ``0`` as a declaration, so a mismatch is + reported against it rather than against a contract. + + Args: + views: The view at this position, one per sample. + position: The position, used in the error message. + + Returns: + One view holding the stacked data and extras. + + Raises: + ValueError: If the views disagree on stream, role or extras. + """ + first = views[0] + for index, view in enumerate(views[1:], start=1): + if (view.stream, view.role) != (first.stream, first.role): + raise ValueError( + f"view {position} is ({first.stream!r}, {first.role!r}) in sample 0 " + f"and ({view.stream!r}, {view.role!r}) in sample {index}" + ) + if set(view.extras) != set(first.extras): + raise ValueError( + f"view {position} has extras {sorted(first.extras)} in sample 0 " + f"and {sorted(view.extras)} in sample {index}" + ) + return View( + data=torch.stack([view.data for view in views]), + stream=first.stream, + role=first.role, + extras={ + key: _stack([view.extras[key] for view in views]) for key in first.extras + }, + ) + + +def _views_and_meta(item: Any) -> tuple[Sequence[View], dict[str, Any]]: + """Splits one dataset item into its views and its per-sample values.""" + if isinstance(item, View): + return [item], {} + if isinstance(item, (list, tuple)) and all(isinstance(x, View) for x in item): + return item, {} + views, rest = item[0], item[1:] + meta: dict[str, Any] = {} + if len(rest) > 0: + meta["target"] = rest[0] + if len(rest) > 1: + meta["filename"] = rest[1] + return views, meta + + +def collate(batch: Sequence[Any]) -> Sample: + """Collates dataset items whose transform returns views. + + Takes no configuration: everything the batch needs to be assembled arrived + with the data. Views are matched by position, and both ``data`` and every + ``extras`` entry are stacked. + + Args: + batch: + The items, each one ``list[View]`` or a tuple whose first element is + ``list[View]``. A second element becomes ``meta["target"]`` and a + third becomes ``meta["filename"]``. + + Returns: + One sample holding the batched views. + + Raises: + ValueError: If the batch is empty or the items disagree on view count. + """ + if len(batch) == 0: + raise ValueError("collate received an empty batch") + + split = [_views_and_meta(item) for item in batch] + views_per_sample = [views for views, _ in split] + + counts = {len(views) for views in views_per_sample} + if len(counts) > 1: + raise ValueError(f"samples in the batch have different view counts: {counts}") + + sample = Sample( + views=[ + _merge(views, position) + for position, views in enumerate(zip(*views_per_sample)) + ] + ) + for key in split[0][1]: + values = [meta[key] for _, meta in split] + sample.meta[key] = ( + _stack(values) + if key != "target" + else ( + torch.stack(values) + if isinstance(values[0], Tensor) + else torch.as_tensor(values) + ) + ) + return sample + + +def legacy_collate(batch: Sequence[Any]) -> tuple[list[Tensor], Tensor, list[str]]: + """Collates into the 1.x ``(views, labels, filenames)`` tuple. + + A shim for training loops written against the old batch type, kept for the + whole 2.x line. + + Args: + batch: The items, as for :func:`collate`. + + Returns: + The views as bare tensors, the labels and the filenames. + """ + sample = collate(batch) + labels = sample.meta.get("target", torch.empty(0, dtype=torch.long)) + filenames = sample.meta.get("filename", []) + return [view.data for view in sample.views], labels, filenames diff --git a/tests/data/test_sample.py b/tests/data/test_sample.py new file mode 100644 index 000000000..ed7bb61c6 --- /dev/null +++ b/tests/data/test_sample.py @@ -0,0 +1,93 @@ +from typing import Any, List, Tuple + +import pytest +import torch + +from lightly.data.sample import Sample, View, collate, legacy_collate + + +def item(target: int = 0, extras: bool = False) -> Tuple[List[View], int]: + views = [ + View(torch.randn(3, 4, 4), extras={"grid": torch.zeros(2)} if extras else {}), + View(torch.randn(3, 4, 4), extras={"grid": torch.ones(2)} if extras else {}), + ] + return views, target + + +def test_selectors_read_the_labels_a_transform_wrote() -> None: + sample = Sample( + views=[ + View(torch.randn(1), role="global"), + View(torch.randn(1), role="local"), + View(torch.randn(1), stream="text", role="local"), + ] + ) + assert len(sample.by_role("global")) == 1 + assert len(sample.by_role("local")) == 2 + assert len(sample.by_stream("text")) == 1 + assert sample.by_role("context") == [] + + +def test_collate_stacks_data_and_keeps_the_view_order() -> None: + sample = collate([item(target=i) for i in range(4)]) + assert [tuple(view.data.shape) for view in sample.views] == [ + (4, 3, 4, 4), + (4, 3, 4, 4), + ] + assert torch.equal(sample.meta["target"], torch.tensor([0, 1, 2, 3])) + + +def test_collate_stacks_every_extras_entry() -> None: + sample = collate([item(extras=True) for _ in range(4)]) + assert tuple(sample.views[1].extras["grid"].shape) == (4, 2) + assert torch.equal(sample.views[1].extras["grid"], torch.ones(4, 2)) + + +def test_collate_takes_views_without_a_target() -> None: + sample = collate([[View(torch.randn(3, 4, 4))] for _ in range(2)]) + assert tuple(sample.views[0].data.shape) == (2, 3, 4, 4) + assert sample.meta == {} + + +def test_collate_keeps_the_filename_of_a_three_tuple() -> None: + batch = [(*item(target=i), f"{i}.jpg") for i in range(2)] + sample = collate(batch) + assert sample.meta["filename"] == ["0.jpg", "1.jpg"] + + +def test_a_ragged_view_count_is_refused() -> None: + batch = [item(), ([View(torch.randn(3, 4, 4))], 0)] + with pytest.raises(ValueError, match="different view counts"): + collate(batch) + + +def test_a_view_labelled_differently_across_samples_is_refused() -> None: + batch = [ + ([View(torch.randn(1), role="global")], 0), + ([View(torch.randn(1), role="local")], 1), + ] + with pytest.raises(ValueError, match="view 0 is"): + collate(batch) + + +def test_extras_that_appear_in_one_sample_only_are_refused() -> None: + batch = [ + ([View(torch.randn(1), extras={"grid": torch.zeros(2)})], 0), + ([View(torch.randn(1))], 1), + ] + with pytest.raises(ValueError, match="extras"): + collate(batch) + + +def test_an_empty_batch_is_refused() -> None: + with pytest.raises(ValueError, match="empty batch"): + collate([]) + + +def test_legacy_collate_yields_the_one_x_tuple() -> None: + views, labels, filenames = legacy_collate( + [(*item(target=i), "a") for i in range(2)] + ) + assert [tuple(view.shape) for view in views] == [(2, 3, 4, 4), (2, 3, 4, 4)] + assert torch.equal(labels, torch.tensor([0, 1])) + assert filenames == ["a", "a"]