add tests

modyn/models/rho_loss_twin_model/rho_loss_twin_model.py

@@ -46,6 +46,8 @@ def set_extra_state(self, state: dict) -> None:
 
     def forward(self, data: torch.Tensor, sample_ids: Optional[list[int]] = None) -> torch.Tensor:
         assert sample_ids is not None
+        # self.training is an internal attribute defined in nn.Module that is updated
+        # whenever .eval() or .train() is called
         if self.training:
             output_tensor = self._training_forward(sample_ids, data)
         else:
@@ -64,8 +66,13 @@ def _eval_forward(self, sample_ids: list[int], data: torch.Tensor) -> torch.Tens
             [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()
-        # when a sample is seen by both models, we route it to model 0
+        # if model 0 did not see any sample, we route all samples to model 0
+        if not seen_by_model0.any():
+            return self._models[0](data)
+        # if model 1 did not see any sample, we route all samples to model 1
+        if not seen_by_model1.any():
+            return self._models[1](data)
+
+        # when a sample is not seen by any model, 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))
+        return torch.where(seen_by_model0.unsqueeze(1), self._models[1](data), self._models[0](data))

modyn/tests/models/test_rho_loss_twin_model.py

@@ -77,17 +77,31 @@ def test_training_forward(mock__eval_forward, current_model, twin_model: RHOLOSS
     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}]
+@pytest.mark.parametrize("exclusive_model", [0, 1])
+@patch.object(RHOLOSSTwinModelModyn, "_training_forward")
+def test_eval_forward_exclusively_route_to_one_model(
+    mock__training_forward, exclusive_model: int, twin_model: RHOLOSSTwinModel
+):
+    twin_model.model._models[1 - exclusive_model].forward = Mock(return_value=torch.zeros(3, 2))
+    twin_model.model._models[exclusive_model].forward = Mock(return_value=torch.ones(3, 2))
+
     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)
+    sample_ids = [1, 2, 3]
+    forward_input = torch.randn(3, 2)
 
+    twin_model.model._models_seen_ids[1 - exclusive_model] = set(sample_ids)
+    twin_model.model._models_seen_ids[exclusive_model] = set()
+
+    output = twin_model.model(forward_input, sample_ids)
+    assert torch.allclose(output, torch.ones(3, 2))
+
+    twin_model.model._models[1 - exclusive_model].forward.assert_not_called()
+    twin_model.model._models[exclusive_model].forward.assert_called_once()
+    # we never call _training_forward in eval mode
+    assert mock__training_forward.call_count == 0
 
-@patch.object(RHOLOSSTwinModelModyn, "_training_forward")
-def test_eval_forward(mock__training_forward, twin_model: RHOLOSSTwinModel):
+
+def test_eval_forward_mixed(twin_model: RHOLOSSTwinModel):
 
     def model0_mock_forward(data: torch.Tensor):
         return torch.zeros(data.shape[0], 10)
@@ -97,13 +111,13 @@ def model1_mock_forward(data: torch.Tensor):
 
     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._models_seen_ids = [{1, 2}, {3, 4}]
     twin_model.model.eval()
 
-    sample_ids = [1, 4, 3, 2]
-    forward_input = torch.randn(4, 2)
+    sample_ids = [1, 4, 3, 2, 5]
+    forward_input = torch.randn(len(sample_ids), 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]))
+    assert torch.allclose(output, torch.tensor([[1.0] * 10, [0.0] * 10, [0.0] * 10, [1.0] * 10, [0.0] * 10]))
 
 
 def test_backup_and_restore_state(twin_model: RHOLOSSTwinModel):