(prototype) Add the View and Sample batch contract

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

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
 

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()

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

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"]