add unit tests

modyn/models/rho_loss_twin_model/rho_loss_twin_model.py

@@ -8,7 +8,6 @@
 logger = logging.getLogger(__name__)
 
 
-
 class RHOLOSSTwinModel:
 
     def __init__(self, model_configuration: dict[str, Any], device: str, amp: bool) -> None:
@@ -26,24 +25,24 @@ def __init__(self, model_configuration: dict[str, Any], device: str, amp: bool)
         rho_model_config = model_configuration["rho_real_model_config"]
         model_handler = getattr(model_module, rho_model_class)
         # we only need the inner model, not the wrapper
-        self.models = [
-            model_handler(rho_model_config, device, amp).model,
-            model_handler(rho_model_config, device, amp).model
-        ]
-        self.models_seen_ids = [
-            set(),
-            set()
-        ]
-        self.current_model = 0
+        self._models = nn.ModuleList(
+            [
+                model_handler(rho_model_config, device, amp).model,
+                model_handler(rho_model_config, device, amp).model,
+            ]
+        )
+        self._models_seen_ids: list[set[int]] = [set(), set()]
+        self._current_model = 0
 
     def get_extra_state(self) -> Any:
         return {
-            "models_seen_ids": self.models_seen_ids,
+            "models_seen_ids": self._models_seen_ids,
         }
 
     def set_extra_state(self, state: dict) -> None:
-        self.models_seen_ids = state["models_seen_ids"]
-        self.current_model = 1
+        self._models_seen_ids = state["models_seen_ids"]
+        # the second time we train on this model, we should switch to the other model
+        self._current_model = 1
 
     def forward(self, data: torch.Tensor, sample_ids: Optional[list[int]] = None) -> torch.Tensor:
         assert sample_ids is not None
@@ -54,13 +53,19 @@ def forward(self, data: torch.Tensor, sample_ids: Optional[list[int]] = None) ->
         return output_tensor
 
     def _training_forward(self, sample_ids: list[int], data: torch.Tensor) -> torch.Tensor:
-        self.models_seen_ids[self.current_model].update(sample_ids)
-        return self.models[self.current_model](data)
+        self._models_seen_ids[self._current_model].update(sample_ids)
+        return self._models[self._current_model](data)
 
     def _eval_forward(self, sample_ids: list[int], data: torch.Tensor) -> torch.Tensor:
-        seen_by_model0 = torch.BoolTensor([sample_id in self.models_seen_ids[0] for sample_id in sample_ids], device=self.device)
-        seen_by_model1 = torch.BoolTensor([sample_id in self.models_seen_ids[1] for sample_id in sample_ids], device=self.device)
+        seen_by_model0 = torch.BoolTensor(
+            [sample_id in self._models_seen_ids[0] for sample_id in sample_ids], device=self.device
+        )
+        seen_by_model1 = torch.BoolTensor(
+            [sample_id in self._models_seen_ids[1] for sample_id in sample_ids], device=self.device
+        )
 
         # assert that a sample is seen by at least one model
         assert (seen_by_model0 | seen_by_model1).all()
-        return torch.where(seen_by_model1, self.models[0](data), self.models[1](data))
+        # when a sample is seen by both models, we route it to model 0
+        # unsqueeze to make seen_by_model1 broadcastable
+        return torch.where(seen_by_model1.unsqueeze(1), self._models[0](data), self._models[1](data))

modyn/tests/models/test_rho_loss_twin_model.py

@@ -0,0 +1,140 @@
+import io
+import tempfile
+from unittest.mock import Mock, call, patch
+
+import pytest
+import torch
+from modyn.models import RHOLOSSTwinModel
+from modyn.models.dummy.dummy import Dummy, DummyModyn
+from modyn.models.rho_loss_twin_model.rho_loss_twin_model import RHOLOSSTwinModelModyn
+
+
+@pytest.fixture
+def twin_model() -> RHOLOSSTwinModel:
+    return RHOLOSSTwinModel(
+        model_configuration={
+            "rho_real_model_class": "Dummy",
+            "rho_real_model_config": {"num_classes": 10},
+        },
+        device="cpu",
+        amp=False,
+    )
+
+
+@patch("modyn.models.Dummy", wraps=Dummy)
+def test_init(MockDummy):
+    model = RHOLOSSTwinModel(
+        model_configuration={
+            "rho_real_model_class": "Dummy",
+            "rho_real_model_config": {"num_classes": 10},
+        },
+        device="cpu",
+        amp=False,
+    )
+    assert len(model.model._models) == 2
+    assert isinstance(model.model._models[0], DummyModyn)
+    assert isinstance(model.model._models[1], DummyModyn)
+    assert model.model._models_seen_ids == [set(), set()]
+    assert model.model._current_model == 0
+    # assert called twice with the same arguments
+    expected_call = call({"num_classes": 10}, "cpu", False)
+    MockDummy.assert_has_calls([expected_call, expected_call])
+
+
+def test_training_forward_missing_sample_ids(twin_model: RHOLOSSTwinModel):
+    with pytest.raises(AssertionError):
+        twin_model.model.train()
+        twin_model.model(torch.randn(3, 2))
+
+
+@pytest.mark.parametrize("current_model", [0, 1])
+@patch.object(RHOLOSSTwinModelModyn, "_eval_forward")
+def test_training_forward(mock__eval_forward, current_model, twin_model: RHOLOSSTwinModel):
+    twin_model.model._models[1 - current_model].forward = Mock()
+    twin_model.model._current_model = current_model
+    assert twin_model.model._models_seen_ids == [set(), set()]
+    twin_model.model.train()
+    sample_ids = [2, 1, 3]
+    forward_input = torch.randn(3, 2)
+    twin_model.model(forward_input, sample_ids)
+    assert twin_model.model._current_model == current_model
+    assert twin_model.model._models_seen_ids[current_model] == set(sample_ids)
+    assert twin_model.model._models_seen_ids[1 - current_model] == set()
+
+    # another forward with the same sample_ids
+    twin_model.model(forward_input, sample_ids)
+    assert twin_model.model._current_model == current_model
+    assert twin_model.model._models_seen_ids[current_model] == set(sample_ids)
+    assert twin_model.model._models_seen_ids[1 - current_model] == set()
+
+    sample_ids = [3, 4]
+    twin_model.model(forward_input, sample_ids)
+    assert twin_model.model._current_model == current_model
+    assert twin_model.model._models_seen_ids[current_model] == {2, 1, 3, 4}
+    assert twin_model.model._models_seen_ids[1 - current_model] == set()
+
+    twin_model.model._models[1 - current_model].forward.assert_not_called()
+    mock__eval_forward.assert_not_called()
+
+
+def test_eval_forward_unseen_sample_id(twin_model: RHOLOSSTwinModel):
+    twin_model.model._models_seen_ids = [{1, 2}, {3, 4}]
+    twin_model.model.eval()
+    sample_ids = [4, 5]
+    forward_input = torch.randn(2, 2)
+    with pytest.raises(AssertionError):
+        twin_model.model(forward_input, sample_ids)
+
+
+@patch.object(RHOLOSSTwinModelModyn, "_training_forward")
+def test_eval_forward(mock__training_forward, twin_model: RHOLOSSTwinModel):
+
+    def model0_mock_forward(data: torch.Tensor):
+        return torch.zeros(data.shape[0], 10)
+
+    def model1_mock_forward(data: torch.Tensor):
+        return torch.ones(data.shape[0], 10)
+
+    twin_model.model._models[0].forward = model0_mock_forward
+    twin_model.model._models[1].forward = model1_mock_forward
+    twin_model.model._models_seen_ids = [{1, 2, 3}, {3, 4}]
+    twin_model.model.eval()
+
+    sample_ids = [1, 4, 3, 2]
+    forward_input = torch.randn(4, 2)
+    output = twin_model.model(forward_input, sample_ids)
+    assert torch.allclose(output, torch.tensor([[1.0] * 10, [0.0] * 10, [0.0] * 10, [1.0] * 10]))
+
+
+def test_backup_and_restore_state(twin_model: RHOLOSSTwinModel):
+    twin_model.model._models_seen_ids = [{1, 2, 3}, {}]
+    assert twin_model.model._current_model == 0
+    twin_model.model._models[0].output.weight = torch.nn.Parameter(torch.zeros(2, 2))
+    twin_model.model._models[0].output.bias = torch.nn.Parameter(torch.zeros(2))
+    twin_model.model._models[1].output.weight = torch.nn.Parameter(torch.ones(2, 2))
+    twin_model.model._models[1].output.bias = torch.nn.Parameter(torch.ones(2))
+
+    with tempfile.NamedTemporaryFile() as model_file:
+        torch.save({"model": twin_model.model.state_dict()}, model_file.name)
+        new_twin_model = RHOLOSSTwinModel(
+            model_configuration={
+                "rho_real_model_class": "Dummy",
+                "rho_real_model_config": {"num_classes": 10},
+            },
+            device="cpu",
+            amp=False,
+        )
+        assert not torch.allclose(new_twin_model.model._models[0].output.weight, torch.zeros(2, 2))
+        assert not torch.allclose(new_twin_model.model._models[0].output.bias, torch.zeros(2))
+        assert not torch.allclose(new_twin_model.model._models[1].output.weight, torch.ones(2, 2))
+        assert not torch.allclose(new_twin_model.model._models[1].output.bias, torch.ones(2))
+        with open(model_file.name, "rb") as f:
+            checkpoint = torch.load(io.BytesIO(f.read()), map_location="cpu")
+            new_twin_model.model.load_state_dict(checkpoint["model"])
+
+        assert new_twin_model.model._models_seen_ids == [{1, 2, 3}, {}]
+        assert new_twin_model.model._current_model == 1
+        assert torch.allclose(new_twin_model.model._models[0].output.weight, torch.zeros(2, 2))
+        assert torch.allclose(new_twin_model.model._models[0].output.bias, torch.zeros(2))
+        assert torch.allclose(new_twin_model.model._models[1].output.weight, torch.ones(2, 2))
+        assert torch.allclose(new_twin_model.model._models[1].output.bias, torch.ones(2))