diff --git a/deeptab/core/serialization.py b/deeptab/core/serialization.py index e6bc60ce..348dff82 100644 --- a/deeptab/core/serialization.py +++ b/deeptab/core/serialization.py @@ -557,12 +557,19 @@ def build_save_bundle( else: task = "regression" if estimator._data_module.regression else "classification" + # set_params() may already have swapped estimator._preprocessor for an + # unfitted one configured for the next fit(); the artifact needs the + # preprocessor the data module was actually fitted with. + fitted_preprocessor = getattr(estimator._data_module, "preprocessor", None) + if fitted_preprocessor is None: + fitted_preprocessor = estimator._preprocessor + artifact_metadata = build_artifact_metadata( estimator=estimator, model_class=type(estimator._estimator), config=estimator.config, data_module=estimator._data_module, - preprocessor=estimator._preprocessor, + preprocessor=fitted_preprocessor, preprocessor_kwargs=getattr(estimator, "_preprocessor_kwargs", {}), task=task, regression=estimator._data_module.regression, @@ -589,7 +596,7 @@ def build_save_bundle( "trainer_config": estimator.trainer_config, "random_state": estimator.random_state, "preprocessor_kwargs": getattr(estimator, "_preprocessor_kwargs", {}), - "preprocessor": estimator._preprocessor, + "preprocessor": fitted_preprocessor, "feature_info": { "num": estimator._data_module.num_feature_info, "cat": estimator._data_module.cat_feature_info, diff --git a/deeptab/models/_mixins/fit.py b/deeptab/models/_mixins/fit.py index 036a0480..11a87daa 100644 --- a/deeptab/models/_mixins/fit.py +++ b/deeptab/models/_mixins/fit.py @@ -346,8 +346,12 @@ def _prepare_continuation( from deeptab.core.default_factories import DefaultDataModuleFactory factory = getattr(self, "_data_module_factory", None) or DefaultDataModuleFactory() + # Continuation reuses the fitted representation, which lives on the + # previous data module; self._preprocessor may already be an unfitted + # replacement created by set_params() for the next full fit. + fitted_preprocessor = getattr(previous, "preprocessor", None) prepared = factory.create( - preprocessor=self._preprocessor, + preprocessor=fitted_preprocessor if fitted_preprocessor is not None else self._preprocessor, batch_size=batch_size, shuffle=shuffle, regression=regression, diff --git a/tests/integration/test_preprocessing_changes_after_fit.py b/tests/integration/test_preprocessing_changes_after_fit.py new file mode 100644 index 00000000..c1af03a5 --- /dev/null +++ b/tests/integration/test_preprocessing_changes_after_fit.py @@ -0,0 +1,45 @@ +"""Changing preprocessing settings on a fitted estimator prepares the *next* fit(). + +The fitted model must keep working: saving it, and continuing to train it with +rebuild=False, both need the preprocessor the data was actually fitted with. +""" + +import numpy as np +import pytest + +from deeptab.configs import MLPConfig, PreprocessingConfig, TrainerConfig +from deeptab.models import MLPLSS, MLPRegressor + +FIT_KWARGS = {"accelerator": "cpu", "logger": False, "enable_progress_bar": False, "enable_model_summary": False} + + +def _fitted(model_class, tmp_path, **fit_options): + rng = np.random.default_rng(0) + features = rng.normal(size=(80, 3)) + model = model_class( + model_config=MLPConfig(layer_sizes=[16]), + preprocessing_config=PreprocessingConfig(), + trainer_config=TrainerConfig(max_epochs=1, checkpoint_path=str(tmp_path)), + random_state=0, + ) + model.fit(features, features[:, 0] * 2, **fit_options, **FIT_KWARGS) + return model, features + + +@pytest.mark.parametrize("model_class, fit_options", [(MLPRegressor, {}), (MLPLSS, {"family": "normal"})]) +def test_save_after_set_params_keeps_the_fitted_preprocessor(model_class, fit_options, tmp_path): + model, features = _fitted(model_class, tmp_path, **fit_options) + expected = model.predict(features) + model.set_params(preprocessing_config__cat_cutoff=0.2) + path = str(tmp_path / "model.deeptab") + model.save(path) + np.testing.assert_allclose(model_class.load(path).predict(features), expected, rtol=1e-6) + + +@pytest.mark.parametrize("model_class, fit_options", [(MLPRegressor, {}), (MLPLSS, {"family": "normal"})]) +def test_continuation_after_set_params_uses_the_fitted_preprocessor(model_class, fit_options, tmp_path): + model, features = _fitted(model_class, tmp_path, **fit_options) + model.set_params(preprocessing_config__cat_cutoff=0.2) + model.fit(features, features[:, 0] * 2, rebuild=False, **fit_options, **FIT_KWARGS) + assert model.is_fitted_ + assert np.isfinite(model.predict(features)).all()