diff --git a/deeptab/training/lightning_module.py b/deeptab/training/lightning_module.py index 9619add..549595a 100644 --- a/deeptab/training/lightning_module.py +++ b/deeptab/training/lightning_module.py @@ -503,7 +503,7 @@ def test_step(self, batch, batch_idx): # type: ignore data, labels = batch if hasattr(self.estimator, "predict_with_candidates") and self.train_features is not None: preds = self.estimator.predict_with_candidates( # type: ignore[reportCallIssue] - *data, candidates_x=self.train_features, candidates_y=self.train_targets + *data, candidate_x=self.train_features, candidate_y=self.train_targets ) else: preds = self(*data) diff --git a/tests/test_candidate_models.py b/tests/test_candidate_models.py new file mode 100644 index 0000000..6dd316f --- /dev/null +++ b/tests/test_candidate_models.py @@ -0,0 +1,68 @@ +"""Regression test: candidate-model step methods must use the right kwarg names. + +``predict_with_candidates`` is defined as ``(self, *data, candidate_x, +candidate_y)`` in TabR and ModernNCA, but ``test_step`` called it with +``candidates_x``/``candidates_y``, so ``trainer.test()`` on those models always +raised TypeError. ``predict_step`` already used the singular names. +""" + +from typing import Any, cast + +import torch +import torch.nn as nn + +from deeptab.configs import MLPConfig +from deeptab.training.lightning_module import TaskModel + + +class _CandidateEstimator(nn.Module): + """Minimal stand-in with TabR's ``predict_with_candidates`` signature.""" + + uses_candidates = True + returns_ensemble = False + + def __init__(self, config=None, feature_information=None, num_classes=1, lss=False, **kwargs): + super().__init__() + self.linear = nn.Linear(4, num_classes) + self.seen = {} + + def forward(self, *data): + return self.linear(data[0][0]) + + def predict_with_candidates(self, *data, candidate_x, candidate_y): + self.seen["candidate_x"] = candidate_x + self.seen["candidate_y"] = candidate_y + return self.forward(*data) + + +def _task_model(): + task = TaskModel( + model_class=_CandidateEstimator, + config=MLPConfig(), + feature_information=({}, {}, {}), + num_classes=1, + ) + task.train_features = ([torch.randn(8, 4)], [], None) + task.train_targets = torch.randn(8, 1) + return task + + +def _batch(batch_size=4): + return (([torch.randn(batch_size, 4)], [], None), torch.randn(batch_size, 1)) + + +def test_test_step_passes_candidate_kwargs(): + task = _task_model() + loss = task.test_step(_batch(), 0) + assert torch.as_tensor(loss).isfinite() + seen = cast(Any, task.estimator).seen + assert seen["candidate_x"] is task.train_features + assert seen["candidate_y"] is task.train_targets + + +def test_predict_step_passes_candidate_kwargs(): + task = _task_model() + # predict_step unpacks the batch as the data tuple itself (no labels). + preds = task.predict_step(_batch()[0], 0) + assert preds.shape == (4, 1) + assert cast(Any, task.estimator).seen["candidate_x"] is task.train_features