feat: Support for RS2 Downsampler

modyn/config/schema/sampling/downsampling_config.py

@@ -3,7 +3,7 @@
 from typing import Annotated, List, Literal, Union
 
 from modyn.config.schema.modyn_base_model import ModynBaseModel
-from pydantic import Field, model_validator
+from pydantic import Field, field_validator, model_validator
 from typing_extensions import Self
 
 
@@ -121,6 +121,33 @@ class RHOLossDownsamplingConfig(BaseDownsamplingConfig):
     il_training_config: ILTrainingConfig = Field(description="The configuration for the IL training.")
 
 
+class RS2DownsamplingConfig(BaseDownsamplingConfig):
+    """Config for the RS2 downsampling strategy."""
+
+    strategy: Literal["RS2"] = "RS2"
+    with_replacement: bool = Field(
+        description=(
+            "Whether we resample from the full TTS each epoch (= True) or train "
+            "on all the data with a different subset each epoch (= False)."
+        )
+    )
+
+    @field_validator("sample_then_batch")
+    @classmethod
+    def sample_then_batch_must_be_true(cls, v: bool) -> bool:
+        if not v:
+            raise ValueError("sample_then_batch must be set to True for this config.")
+        return v
+
+    @field_validator("period")
+    @classmethod
+    def only_support_period_one(cls, v: int) -> int:
+        if v != 0:
+            # RS2 requires us to resample every epoch.
+            raise ValueError("period must be set to 1 for this config.")
+        return v
+
+
 SingleDownsamplingConfig = Annotated[
     Union[
         UncertaintyDownsamplingConfig,

modyn/tests/trainer_server/internal/trainer/remote_downsamplers/test_rs2_downsampling.py

@@ -0,0 +1,172 @@
+import torch
+from modyn.trainer_server.internal.trainer.remote_downsamplers.abstract_remote_downsampling_strategy import (
+    get_tensors_subset,
+)
+from modyn.trainer_server.internal.trainer.remote_downsamplers.remote_rs2_downsampling import RemoteRS2Downsampling
+
+
+def test_init():
+    pipeline_id = 0
+    trigger_id = 0
+    batch_size = 32
+    params_from_selector = {"replacement": True, "downsampling_ratio": 50}
+    per_sample_loss = None
+    device = "cpu"
+
+    downsampler = RemoteRS2Downsampling(
+        pipeline_id, trigger_id, batch_size, params_from_selector, per_sample_loss, device
+    )
+
+    assert downsampler.pipeline_id == pipeline_id
+    assert downsampler.trigger_id == trigger_id
+    assert downsampler.batch_size == batch_size
+    assert downsampler.device == device
+    assert not downsampler.forward_required
+    assert not downsampler.supports_bts
+    assert downsampler._all_sample_ids == []
+    assert downsampler._subsets == []
+    assert downsampler._current_subset == -1
+    assert downsampler._with_replacement == params_from_selector["replacement"]
+    assert downsampler._max_subset == -1
+    assert downsampler._first_epoch
+
+
+def test_inform_samples():
+    pipeline_id = 0
+    trigger_id = 0
+    batch_size = 32
+    params_from_selector = {"replacement": True, "downsampling_ratio": 50}
+    per_sample_loss = None
+    device = "cpu"
+
+    downsampler = RemoteRS2Downsampling(
+        pipeline_id, trigger_id, batch_size, params_from_selector, per_sample_loss, device
+    )
+
+    sample_ids = [1, 2, 3, 4, 5]
+    forward_output = torch.randn(5, 10)
+    target = torch.randint(0, 10, (5,))
+
+    downsampler.inform_samples(sample_ids, forward_output, target)
+
+    assert downsampler._all_sample_ids == sample_ids
+    downsampler.inform_samples(sample_ids, forward_output, target)
+    assert downsampler._all_sample_ids == 2 * sample_ids
+    # Now it should not change anymore
+    downsampler.select_points()
+    downsampler.inform_samples(sample_ids, forward_output, target)
+    assert set(downsampler._all_sample_ids) == set(sample_ids)
+    assert len(downsampler._all_sample_ids) == 2 * len(sample_ids)
+
+
+def test_multiple_epochs_with_replacement():
+    pipeline_id = 0
+    trigger_id = 0
+    batch_size = 32
+    params_from_selector = {"replacement": True, "downsampling_ratio": 50}
+    per_sample_loss = None
+    device = "cpu"
+
+    downsampler = RemoteRS2Downsampling(
+        pipeline_id, trigger_id, batch_size, params_from_selector, per_sample_loss, device
+    )
+    with torch.inference_mode(mode=(not downsampler.requires_grad)):
+        sample_ids = list(range(10))
+        data = torch.randn(10, 10)
+        target = torch.randint(0, 10, (10,))
+
+        for _ in range(3):
+            downsampler.inform_samples(sample_ids, data, target)
+            selected_ids, weights = downsampler.select_points()
+            sampled_data, sampled_target = get_tensors_subset(selected_ids, data, target, sample_ids)
+
+            assert len(set(selected_ids)) == 5
+            assert weights.shape == (5,)
+            assert all(idx in sample_ids for idx in selected_ids)
+            assert sampled_data.shape == (5, 10)
+            assert sampled_target.shape == (5,)
+
+
+def test_multiple_epochs_without_replacement():
+    pipeline_id = 0
+    trigger_id = 0
+    batch_size = 32
+    params_from_selector = {"replacement": False, "downsampling_ratio": 50}
+    per_sample_loss = None
+    device = "cpu"
+
+    downsampler = RemoteRS2Downsampling(
+        pipeline_id, trigger_id, batch_size, params_from_selector, per_sample_loss, device
+    )
+    with torch.inference_mode(mode=(not downsampler.requires_grad)):
+
+        sample_ids = list(range(10))
+        data = torch.randn(10, 10)
+        target = torch.randint(0, 10, (10,))
+
+        # Epoch 1
+        downsampler.inform_samples(sample_ids, data, target)
+        epoch1_ids, weights = downsampler.select_points()
+        sampled_data, sampled_target = get_tensors_subset(epoch1_ids, data, target, sample_ids)
+
+        assert len(set(epoch1_ids)) == 5
+        assert weights.shape == (5,)
+        assert all(idx in sample_ids for idx in epoch1_ids)
+        assert sampled_data.shape == (5, 10)
+        assert sampled_target.shape == (5,)
+
+        # Epoch 2
+        downsampler.inform_samples(sample_ids, data, target)
+        epoch2_ids, weights = downsampler.select_points()
+        sampled_data, sampled_target = get_tensors_subset(epoch2_ids, data, target, sample_ids)
+
+        assert len(set(epoch2_ids)) == 5
+        assert weights.shape == (5,)
+        assert all(idx in sample_ids for idx in epoch2_ids)
+        assert not any(idx in epoch1_ids for idx in epoch2_ids)  # No overlap across epochs
+        assert sampled_data.shape == (5, 10)
+        assert sampled_target.shape == (5,)
+
+        # Epoch 3
+        downsampler.inform_samples(sample_ids, data, target)
+        epoch3_ids, weights = downsampler.select_points()
+        sampled_data, sampled_target = get_tensors_subset(epoch3_ids, data, target, sample_ids)
+
+        assert len(set(epoch3_ids)) == 5
+        assert weights.shape == (5,)
+        assert all(idx in sample_ids for idx in epoch3_ids)
+        assert all(idx in epoch1_ids or idx in epoch2_ids for idx in epoch3_ids)  # There needs to be overlap now
+        # but (with very high probability, this might be flaky lets see) there is some difference
+        assert any(idx not in epoch1_ids for idx in epoch3_ids)
+        assert sampled_data.shape == (5, 10)
+        assert sampled_target.shape == (5,)
+
+
+def test_multiple_epochs_without_replacement_leftover_data():
+    pipeline_id = 0
+    trigger_id = 0
+    batch_size = 32
+    params_from_selector = {"replacement": False, "downsampling_ratio": 40}
+    per_sample_loss = None
+    device = "cpu"
+
+    downsampler = RemoteRS2Downsampling(
+        pipeline_id, trigger_id, batch_size, params_from_selector, per_sample_loss, device
+    )
+    with torch.inference_mode(mode=(not downsampler.requires_grad)):
+        sample_ids = list(range(10))
+        data = torch.randn(10, 10)
+        target = torch.randint(0, 10, (10,))
+
+        for _ in range(3):
+            downsampler.inform_samples(sample_ids, data, target)
+
+            selected_ids, weights = downsampler.select_points()
+            sampled_data, sampled_target = get_tensors_subset(selected_ids, data, target, sample_ids)
+            assert len(set(selected_ids)) == 4
+            assert weights.shape == (4,)
+            assert sampled_data.shape == (4, 10)
+            assert sampled_target.shape == (4,)
+
+            assert all(idx in sample_ids for idx in selected_ids)
+            assert len(set(selected_ids)) == len(selected_ids)

modyn/trainer_server/internal/trainer/pytorch_trainer.py

@@ -374,6 +374,8 @@ def train(self) -> None:  # pylint: disable=too-many-locals, too-many-branches
             self._info(f"Training will stop when the number of samples to pass reaches {self.num_samples_to_pass}.")
 
         if self._downsampling_mode == DownsamplingMode.BATCH_THEN_SAMPLE:
+            # assertion since model validation by pydantic should catch this.
+            assert self._downsampler.supports_bts, "The downsampler does not support batch then sample"
             # We cannot pass the target size from the trainer server since that depends on StB vs BtS.
             post_downsampling_size = max(int(self._downsampler.downsampling_ratio * self._batch_size / 100), 1)
             assert post_downsampling_size < self._batch_size
@@ -692,7 +694,7 @@ def downsample_batch(
         self.start_embedding_recording_if_needed()
 
         with torch.inference_mode(mode=(not self._downsampler.requires_grad)):
-            big_batch_output = self._model.model(data)
+            big_batch_output = self._model.model(data) if self._downsampler.forward_required else torch.Tensor()
             embeddings = self.get_embeddings_if_recorded()
             self._downsampler.inform_samples(sample_ids, big_batch_output, target, embeddings)
 
@@ -831,7 +833,7 @@ def _iterate_dataloader_and_compute_scores(
             with torch.inference_mode(mode=(not self._downsampler.requires_grad)):
                 with torch.autocast(self._device_type, enabled=self._amp):
                     # compute the scores and accumulate them
-                    model_output = self._model.model(data)
+                    model_output = self._model.model(data) if self._downsampler.forward_required else torch.Tensor()
                     embeddings = self.get_embeddings_if_recorded()
                     self._downsampler.inform_samples(sample_ids, model_output, target, embeddings)
 

modyn/trainer_server/internal/trainer/remote_downsamplers/abstract_remote_downsampling_strategy.py

@@ -62,6 +62,16 @@ def __init__(
         # CoresetSupportingModule for model implementations.
         self.requires_coreset_supporting_module = False
 
+        # Some methods might not need information from forward pass (e.g. completely random)
+        # Most do (definition), hence we default to True
+        # We might want to refactor those downsamplers to presamplers and support some
+        # adaptivity at the selector, but for now we allow random downsamplers mostly
+        # to support RS2.
+        self.forward_required = True
+
+        # Some methods might only support StB, not BtS.
+        self.supports_bts = True
+
     @abstractmethod
     def init_downsampler(self) -> None:
         raise NotImplementedError

modyn/trainer_server/internal/trainer/remote_downsamplers/remote_rs2_downsampling.py

@@ -0,0 +1,81 @@
+import logging
+import random
+from typing import Any, Optional
+
+import torch
+from modyn.trainer_server.internal.trainer.remote_downsamplers.abstract_remote_downsampling_strategy import (
+    AbstractRemoteDownsamplingStrategy,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class RemoteRS2Downsampling(AbstractRemoteDownsamplingStrategy):
+    """
+    Method adapted from REPEATED RANDOM SAMPLING FOR MINIMIZING THE TIME-TO-ACCURACY OF LEARNING (Okanovic+, 2024)
+    https://openreview.net/pdf?id=JnRStoIuTe
+    """
+
+    def __init__(
+        self,
+        pipeline_id: int,
+        trigger_id: int,
+        batch_size: int,
+        params_from_selector: dict,
+        per_sample_loss: Any,
+        device: str,
+    ) -> None:
+        super().__init__(pipeline_id, trigger_id, batch_size, params_from_selector, device)
+        self.forward_required = False
+        self.supports_bts = False
+        self._all_sample_ids: list[int] = []
+        self._subsets: list[list[int]] = []
+        self._current_subset = -1
+        self._with_replacement: bool = params_from_selector["replacement"]
+        self._max_subset = -1
+        self._first_epoch = True
+
+    def init_downsampler(self) -> None:
+        pass  # We take care of that in inform_samples
+
+    def inform_samples(
+        self,
+        sample_ids: list[int],
+        forward_output: torch.Tensor,
+        target: torch.Tensor,
+        embedding: Optional[torch.Tensor] = None,
+    ) -> None:
+        # We only need to collect the sample information once
+        if self._first_epoch:
+            self._all_sample_ids.extend(sample_ids)
+
+    def _epoch_step_wr(self, target_size: int) -> None:
+        self._subsets = [self._all_sample_ids[:target_size]]
+        self._current_subset = 0
+
+    def _epoch_step_r(self, target_size: int) -> None:
+        self._max_subset = len(self._all_sample_ids) // target_size
+        self._current_subset += 1
+        if self._current_subset >= self._max_subset or len(self._subsets) == 0:
+            self._current_subset = 0
+            self._subsets = [
+                self._all_sample_ids[i * target_size : (i + 1) * target_size] for i in range(self._max_subset)
+            ]
+
+    def _epoch_step(self) -> None:
+        target_size = max(int(self.downsampling_ratio * len(self._all_sample_ids) / 100), 1)
+        random.shuffle(self._all_sample_ids)
+
+        if self._with_replacement:
+            self._epoch_step_wr(target_size)
+        else:
+            self._epoch_step_r(target_size)
+
+    def select_points(self) -> tuple[list[int], torch.Tensor]:
+        self._first_epoch = False
+        self._epoch_step()
+        return self._subsets[self._current_subset], torch.ones(len(self._subsets[self._current_subset]))
+
+    @property
+    def requires_grad(self) -> bool:
+        return False