Adding StratifiedBatchSampler for guaranteed subgroup representation

docs/source/distributed.rst

@@ -72,6 +72,7 @@ ignite.distributed.auto
     :toctree: generated
 
     DistributedProxySampler
+    StratifiedBatchSampler
     auto_dataloader
     auto_model
     auto_optim

ignite/distributed/auto.py

@@ -9,15 +9,15 @@
 from torch.optim.optimizer import Optimizer
 from torch.utils.data import DataLoader, Dataset, IterableDataset
 from torch.utils.data.distributed import DistributedSampler
-from torch.utils.data.sampler import Sampler
+from torch.utils.data.sampler import BatchSampler, Sampler
 
 from ignite.distributed import utils as idist
 from ignite.distributed.comp_models import horovod as idist_hvd
 from ignite.distributed.comp_models import native as idist_native
 from ignite.distributed.comp_models import xla as idist_xla
 from ignite.utils import setup_logger
 
-__all__ = ["auto_dataloader", "auto_model", "auto_optim", "DistributedProxySampler"]
+__all__ = ["auto_dataloader", "auto_model", "auto_optim", "DistributedProxySampler", "StratifiedBatchSampler"]
 
 
 def auto_dataloader(dataset: Dataset, **kwargs: Any) -> DataLoader | _MpDeviceLoader:
@@ -330,6 +330,160 @@ def __iter__(self) -> Iterator:
         return iter(indices)
 
 
+class StratifiedBatchSampler(BatchSampler):
+    """Batch sampler guaranteeing that every group is represented in every batch.
+
+    Unlike `torch WeightedRandomSampler`_, which only makes rare groups more likely, this sampler puts a fixed
+    number of samples of each group into each batch. This is required by methods where a missing group zeroes out
+    a loss term, e.g. Group Distributionally Robust Optimization or contrastive learning.
+
+    It is a batch sampler and should be passed to ``batch_sampler`` of a dataloader, not to ``sampler``:
+
+    .. code-block:: python
+
+        batch_sampler = StratifiedBatchSampler(group_ids, batch_size=32)
+        dataloader = DataLoader(dataset, batch_sampler=batch_sampler)
+
+        @trainer.on(Events.EPOCH_STARTED)
+        def set_epoch(engine):
+            batch_sampler.set_epoch(engine.state.epoch)
+
+    Args:
+        labels: group id of every dataset sample, as a sequence, numpy array or 1-D tensor. Groups are derived from
+            the distinct values, which can be of any hashable type.
+        batch_size: number of samples per batch. Should be larger or equal to the number of groups. Every yielded
+            batch has exactly this size. If it is not divisible by the number of groups, the leftover slots are
+            given to a rotating subset of the groups.
+        strategy: if ``"oversample"`` (default), an epoch lasts until the largest group is exhausted and smaller
+            groups are cycled, such that no sample is discarded. If ``"undersample"``, an epoch stops once the
+            smallest group is exhausted and larger groups are truncated. Groups are reshuffled at every epoch, so
+            the truncated part of a larger group is different at every epoch and no sample is permanently unused.
+        num_replicas: number of processes participating in distributed training. Defaults to
+            :meth:`~ignite.distributed.utils.get_world_size`.
+        rank: rank of the current process within ``num_replicas``. Defaults to
+            :meth:`~ignite.distributed.utils.get_rank`.
+        seed: random seed used to shuffle the groups. The seed of an epoch is ``seed + epoch``.
+
+    .. note::
+        In a distributed configuration, groups are sharded across ranks individually, so that every batch of every
+        rank still contains every group. Wrapping this sampler with
+        :class:`~ignite.distributed.auto.DistributedProxySampler` or `torch DistributedSampler`_ instead would shard
+        the flat index list and could leave a rank without any sample of a rare group. Each group should therefore
+        have at least
+        ``num_replicas`` samples. Groups are truncated to a multiple of ``num_replicas``, so that all ranks yield
+        the same number of batches.
+
+    .. note::
+        As for `torch DistributedSampler`_, :meth:`set_epoch` should be called at the beginning of every epoch,
+        otherwise the same batches are produced at every epoch.
+
+    .. _torch WeightedRandomSampler:
+        https://pytorch.org/docs/stable/data.html#torch.utils.data.WeightedRandomSampler
+    .. _torch DistributedSampler:
+        https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
+
+    .. versionadded:: 0.6.0
+    """
+
+    def __init__(
+        self,
+        labels: Any,
+        batch_size: int,
+        strategy: str = "oversample",
+        num_replicas: int | None = None,
+        rank: int | None = None,
+        seed: int = 0,
+    ) -> None:
+        if strategy not in ("oversample", "undersample"):
+            raise ValueError(f"Argument strategy should be 'oversample' or 'undersample', but given {strategy}")
+
+        ndim = getattr(labels, "ndim", 1)
+        if ndim != 1:
+            raise ValueError(f"Argument labels should be one-dimensional, but given {ndim} dimensions")
+        if hasattr(labels, "tolist"):
+            labels = labels.tolist()
+
+        groups: dict[Any, list[int]] = {}
+        for index, label in enumerate(labels):
+            groups.setdefault(label, []).append(index)
+        if len(groups) == 0:
+            raise ValueError("Argument labels should not be empty")
+
+        self.groups = list(groups.values())
+        num_groups = len(self.groups)
+        if batch_size < num_groups:
+            raise ValueError(
+                f"Argument batch_size should be larger or equal to the number of groups ({num_groups}), "
+                f"but given {batch_size}"
+            )
+
+        self.num_replicas = idist.get_world_size() if num_replicas is None else num_replicas
+        self.rank = idist.get_rank() if rank is None else rank
+        if not 0 <= self.rank < self.num_replicas:
+            raise ValueError(f"Argument rank should be in [0, {self.num_replicas}), but given {self.rank}")
+
+        # all ranks drop the same tail, otherwise they yield a different number of batches and collective ops hang
+        self.shard_sizes = [len(group) // self.num_replicas for group in self.groups]
+        if min(self.shard_sizes) == 0:
+            smallest = min(len(group) for group in self.groups)
+            raise ValueError(
+                f"Smallest group has {smallest} samples, less than num_replicas={self.num_replicas}. "
+                "Every group should have at least one sample per rank"
+            )
+
+        self.samples_per_group = batch_size // num_groups
+        self.leftover_slots = batch_size % num_groups
+        if strategy == "oversample":
+            self.num_batches = max(-(-size // self.samples_per_group) for size in self.shard_sizes)
+        else:
+            self.num_batches = max(1, min(size // self.samples_per_group for size in self.shard_sizes))
+
+        self.strategy = strategy
+        self.seed = seed
+        self.epoch = 0
+        super().__init__(sampler=None, batch_size=batch_size, drop_last=False)  # type: ignore[arg-type]
+
+    def set_epoch(self, epoch: int) -> None:
+        """Set the epoch used to seed the shuffling of the groups.
+
+        Args:
+            epoch: epoch number.
+        """
+        self.epoch = epoch
+
+    def _cycle(self, shard: list[int], generator: torch.Generator) -> Iterator[int]:
+        # reshuffled on every pass, unlike itertools.cycle which replays a frozen order
+        while True:
+            for i in torch.randperm(len(shard), generator=generator).tolist():
+                yield shard[i]
+
+    def __iter__(self) -> Iterator[list[int]]:
+        generator = torch.Generator().manual_seed(self.seed + self.epoch)
+
+        cycles = []
+        for group, size in zip(self.groups, self.shard_sizes):
+            perm = torch.randperm(len(group), generator=generator).tolist()
+            shuffled = [group[i] for i in perm]
+            # same permutation on every rank, so that the strided shards are disjoint
+            cycles.append(self._cycle(shuffled[self.rank :: self.num_replicas][:size], generator))
+
+        offset = 0
+        for _ in range(self.num_batches):
+            # rotated, such that no group systematically gets the slots left over by the integer division
+            bonus = {(offset + i) % len(cycles) for i in range(self.leftover_slots)}
+            offset += self.leftover_slots
+
+            batch = []
+            for i, cycle in enumerate(cycles):
+                for _ in range(self.samples_per_group + (1 if i in bonus else 0)):
+                    batch.append(next(cycle))
+
+            yield [batch[i] for i in torch.randperm(len(batch), generator=generator).tolist()]
+
+    def __len__(self) -> int:
+        return self.num_batches
+
+
 if idist.has_xla_support:
     import torch_xla.core.xla_model as xm
     from torch_xla.distributed.parallel_loader import ParallelLoader

tests/ignite/distributed/test_auto.py

@@ -1,4 +1,5 @@
 import os
+from collections import Counter
 
 import pytest
 import torch
@@ -11,7 +12,14 @@
 from torch.utils.data.sampler import BatchSampler, RandomSampler, Sampler, SequentialSampler, WeightedRandomSampler
 
 import ignite.distributed as idist
-from ignite.distributed.auto import auto_dataloader, auto_model, auto_optim, DistributedProxySampler
+from ignite.distributed.auto import (
+    auto_dataloader,
+    auto_model,
+    auto_optim,
+    DistributedProxySampler,
+    StratifiedBatchSampler,
+)
+from ignite.engine.deterministic import ReproducibleBatchSampler
 from tests.ignite import is_mps_available_and_functional
 
 
@@ -315,3 +323,125 @@ def test_dist_proxy_sampler():
 
     with pytest.raises(TypeError, match=r"Argument sampler must not be a distributed sampler already"):
         DistributedProxySampler(DistributedSampler(sampler, num_replicas=num_replicas, rank=0))
+
+
+_STRATIFIED_LABELS = [0] * 100 + [1] * 20 + [2] * 5
+
+
+def _groups_of(batch, labels=_STRATIFIED_LABELS):
+    return Counter(labels[i] for i in batch)
+
+
+def test_stratified_batch_sampler():
+    sampler = StratifiedBatchSampler(_STRATIFIED_LABELS, batch_size=9)
+    for batch in sampler:
+        assert len(batch) == 9
+        assert _groups_of(batch) == Counter({0: 3, 1: 3, 2: 3})
+
+    # batch_size is not divisible by the number of groups: leftover slots rotate over the groups
+    sampler = StratifiedBatchSampler(_STRATIFIED_LABELS, batch_size=10)
+    counts = Counter()
+    for batch in sampler:
+        assert len(batch) == 10
+        assert set(_groups_of(batch)) == {0, 1, 2}
+        counts.update(_groups_of(batch))
+    assert max(counts.values()) - min(counts.values()) <= len(sampler)
+
+    labels = ["cat", "dog", "cat", "dog", "cat", "dog"]
+    sampler = StratifiedBatchSampler(labels, batch_size=2)
+    for batch in sampler:
+        assert {labels[i] for i in batch} == {"cat", "dog"}
+
+    sampler = StratifiedBatchSampler(torch.tensor(_STRATIFIED_LABELS), batch_size=9)
+    assert len(list(sampler)) == len(sampler)
+
+
+def test_stratified_batch_sampler_strategies():
+    oversample = StratifiedBatchSampler(_STRATIFIED_LABELS, batch_size=3, strategy="oversample")
+    undersample = StratifiedBatchSampler(_STRATIFIED_LABELS, batch_size=3, strategy="undersample")
+
+    # oversample lasts until the largest group is exhausted, undersample until the smallest one is
+    assert len(oversample) == 100
+    assert len(undersample) == 5
+    assert len(list(oversample)) == len(oversample)
+    assert len(list(undersample)) == len(undersample)
+
+    indices = [i for batch in undersample for i in batch]
+    assert len(indices) == len(set(indices))
+
+    # the discarded part of the largest group changes at every epoch
+    undersample.set_epoch(0)
+    first = {i for batch in undersample for i in batch if _STRATIFIED_LABELS[i] == 0}
+    undersample.set_epoch(1)
+    second = {i for batch in undersample for i in batch if _STRATIFIED_LABELS[i] == 0}
+    assert first != second
+
+    oversample.set_epoch(0)
+    seen = Counter(i for batch in oversample for i in batch)
+    assert set(seen) == set(range(len(_STRATIFIED_LABELS)))
+
+
+def test_stratified_batch_sampler_seed():
+    sampler = StratifiedBatchSampler(_STRATIFIED_LABELS, batch_size=9, seed=42)
+    batches = list(sampler)
+    assert list(sampler) == batches
+
+    sampler.set_epoch(1)
+    assert list(sampler) != batches
+
+    other = StratifiedBatchSampler(_STRATIFIED_LABELS, batch_size=9, seed=43)
+    assert list(other) != batches
+
+
+def test_stratified_batch_sampler_num_replicas():
+    num_replicas = 4
+    samplers = [
+        StratifiedBatchSampler(_STRATIFIED_LABELS, batch_size=9, num_replicas=num_replicas, rank=i)
+        for i in range(num_replicas)
+    ]
+
+    # all ranks should yield the same number of batches, otherwise collective ops hang
+    assert len({len(s) for s in samplers}) == 1
+
+    indices_per_rank = []
+    for s in samplers:
+        s.set_epoch(0)
+        rank_indices = []
+        for batch in s:
+            assert set(_groups_of(batch)) == {0, 1, 2}
+            rank_indices += batch
+        indices_per_rank.append(set(rank_indices))
+
+    for i in range(num_replicas):
+        for j in range(i + 1, num_replicas):
+            assert not indices_per_rank[i] & indices_per_rank[j]
+
+
+def test_stratified_batch_sampler_with_dataloader():
+    dataloader = DataLoader(DummyDS(length=125), batch_sampler=StratifiedBatchSampler(_STRATIFIED_LABELS, 9))
+    for batch in dataloader:
+        assert set(_groups_of(batch.tolist())) == {0, 1, 2}
+
+    # a deterministic engine replaces the batch sampler, which requires a torch BatchSampler
+    sampler = StratifiedBatchSampler(_STRATIFIED_LABELS, batch_size=9)
+    assert len(list(ReproducibleBatchSampler(sampler))) == len(sampler)
+
+
+def test_stratified_batch_sampler_asserts():
+    with pytest.raises(ValueError, match=r"Argument strategy should be 'oversample' or 'undersample'"):
+        StratifiedBatchSampler(_STRATIFIED_LABELS, batch_size=9, strategy="drop")
+
+    with pytest.raises(ValueError, match=r"Argument labels should be one-dimensional"):
+        StratifiedBatchSampler(torch.zeros(4, 2), batch_size=2)
+
+    with pytest.raises(ValueError, match=r"Argument labels should not be empty"):
+        StratifiedBatchSampler([], batch_size=1)
+
+    with pytest.raises(ValueError, match=r"Argument batch_size should be larger or equal to the number of groups"):
+        StratifiedBatchSampler(_STRATIFIED_LABELS, batch_size=2)
+
+    with pytest.raises(ValueError, match=r"Argument rank should be in \[0, 2\)"):
+        StratifiedBatchSampler(_STRATIFIED_LABELS, batch_size=9, num_replicas=2, rank=2)
+
+    with pytest.raises(ValueError, match=r"Smallest group has 5 samples, less than num_replicas=8"):
+        StratifiedBatchSampler(_STRATIFIED_LABELS, batch_size=9, num_replicas=8, rank=0)