diff --git a/deeptab/models/_mixins/predict.py b/deeptab/models/_mixins/predict.py index 464f353..08fbe80 100644 --- a/deeptab/models/_mixins/predict.py +++ b/deeptab/models/_mixins/predict.py @@ -162,9 +162,18 @@ def encode(self, X, embeddings=None, batch_size=64): encoded_dataset = self._data_module.preprocess_new_data(X, embeddings) data_loader = DataLoader(encoded_dataset, batch_size=batch_size, shuffle=False) - encoded_outputs = [] - for batch in tqdm(data_loader): - emb = self._task_model.estimator.encode(batch) # type: ignore[union-attr] - encoded_outputs.append(emb) + # Embeddings must be deterministic: without eval() dropout stays active + # after fit() and two calls on the same rows disagree. + was_training = self._task_model.training # type: ignore[attr-defined] + self._task_model.eval() + try: + encoded_outputs = [] + with torch.no_grad(): + for batch in tqdm(data_loader): + emb = self._task_model.estimator.encode(batch) # type: ignore[union-attr] + encoded_outputs.append(emb) + finally: + if was_training: + self._task_model.train() return torch.cat(encoded_outputs, dim=0) diff --git a/deeptab/models/lss_base.py b/deeptab/models/lss_base.py index 16a8f58..9925657 100644 --- a/deeptab/models/lss_base.py +++ b/deeptab/models/lss_base.py @@ -505,11 +505,20 @@ def encode(self, X, batch_size=64): data_loader = DataLoader(encoded_dataset, batch_size=batch_size, shuffle=False) - # Process data in batches + # Process data in batches. preprocess_new_data yields + # (num_features, cat_features, embeddings) triples, and encode() takes the + # whole batch -- matching the base _PredictMixin.encode implementation. + was_training = self._task_model.training # type: ignore[union-attr] + self._task_model.eval() # type: ignore[union-attr] encoded_outputs = [] - for num_features, cat_features in tqdm(data_loader): - embeddings = self._task_model.estimator.encode(num_features, cat_features) # type: ignore[union-attr] # Call your encode function - encoded_outputs.append(embeddings) + try: + with torch.no_grad(): + for batch in tqdm(data_loader): + embeddings = self._task_model.estimator.encode(batch) # type: ignore[union-attr] + encoded_outputs.append(embeddings) + finally: + if was_training: + self._task_model.train() # type: ignore[union-attr] # Concatenate all encoded outputs encoded_outputs = torch.cat(encoded_outputs, dim=0) diff --git a/tests/test_encode.py b/tests/test_encode.py new file mode 100644 index 0000000..57d13c1 --- /dev/null +++ b/tests/test_encode.py @@ -0,0 +1,61 @@ +"""Regression tests for the public encode() API. + +encode() never switched the model to eval mode, so embeddings taken after fit() +were dropout-corrupted and differed between calls. The LSS override additionally +unpacked 2 values from the 3-element batches preprocess_new_data yields, so it +always raised ValueError. +""" + +import numpy as np +import pandas as pd + +from deeptab.configs import FTTransformerConfig +from deeptab.models import FTTransformerLSS, FTTransformerRegressor + +QUIET = { + "accelerator": "cpu", + "devices": 1, + "enable_progress_bar": False, + "enable_model_summary": False, + "logger": False, +} + + +def _data(n=48, seed=0): + rng = np.random.default_rng(seed) + X = pd.DataFrame({"a": rng.normal(size=n), "b": rng.normal(size=n), "g": rng.choice(list("xyz"), n)}) + return X, rng.normal(size=n) + + +def _config(dropout=0.5): + return FTTransformerConfig(d_model=16, n_layers=1, n_heads=2, attn_dropout=dropout) + + +def test_encode_is_deterministic_after_fit(): + """Dropout must be disabled: two calls on the same rows must agree.""" + X, y = _data() + model = FTTransformerRegressor(model_config=_config()) + model.fit(X, y, max_epochs=1, batch_size=16, **QUIET) + + first = np.asarray(model.encode(X)) + second = np.asarray(model.encode(X)) + np.testing.assert_allclose(first, second, rtol=1e-5, atol=1e-6) + + +def test_encode_restores_training_mode(): + X, y = _data() + model = FTTransformerRegressor(model_config=_config()) + model.fit(X, y, max_epochs=1, batch_size=16, **QUIET) + + model._task_model.train() # type: ignore[union-attr] + model.encode(X) + assert model._task_model.training is True # type: ignore[union-attr] + + +def test_lss_encode_returns_embeddings(): + X, y = _data() + model = FTTransformerLSS(model_config=_config(dropout=0.0)) + model.fit(X, y, family="normal", max_epochs=1, batch_size=16, **QUIET) + + encoded = np.asarray(model.encode(X)) + assert encoded.shape[0] == len(X)