diff --git a/src/mother/ml/models/m_tabpfn.py b/src/mother/ml/models/m_tabpfn.py index eb872d7..3fc24f7 100644 --- a/src/mother/ml/models/m_tabpfn.py +++ b/src/mother/ml/models/m_tabpfn.py @@ -536,6 +536,11 @@ def _set_random_state(self) -> None: torch.manual_seed(self.random_state) np.random.seed(self.random_state) + def _prepare_prefitted_model_for_embeddings(self) -> None: + if self.model is None: + raise RuntimeError("A pre-fitted model is required to extract embeddings.") + self.model.use_autocast_ = False + def _get_best_embeddings( self, embeddings: np.ndarray, @@ -633,6 +638,7 @@ def fit( module_logger.info( "A pre-fitted model has been given. The new data will not be used for fitting the model." ) + self._prepare_prefitted_model_for_embeddings() self.train_embeddings_ = self.model.get_embeddings(X_array) self._embedding_dim = self.train_embeddings_.shape[1] else: @@ -777,6 +783,8 @@ def transform( X_array = np.asarray(X, dtype=np.float32) # Get embeddings for new data using the main model + if self.pre_fitted: + self._prepare_prefitted_model_for_embeddings() embeddings = self.model.get_embeddings(X_array) # collapse the additional column caused by estimators (avg) if len(embeddings.shape) == 3: diff --git a/src/mother/optimization/core.py b/src/mother/optimization/core.py index befef5e..8aaefbb 100644 --- a/src/mother/optimization/core.py +++ b/src/mother/optimization/core.py @@ -13,7 +13,11 @@ import sklearn.metrics as skl_metrics import sklearn.model_selection as skl_model_sel from optuna.study import Study, StudyDirection -from optuna.terminator import TerminatorCallback, report_cross_validation_scores +from optuna.terminator import ( + Terminator, + TerminatorCallback, + report_cross_validation_scores, +) from sklearn.pipeline import Pipeline from mother import utils as mother_utils @@ -133,7 +137,7 @@ def __init__( tuning_direction: typing.Union[StudyDirection, str] = StudyDirection.MAXIMIZE, n_trials_optuna: int = 100, n_threads_optuna: int = 1, - n_startup_trials: int = 12, + n_startup_trials: int = 20, seed: int = 42, **kwargs, ): @@ -157,17 +161,22 @@ def __init__( self.study: typing.Optional[Study] = None - def get_callbacks(self): + def get_callbacks(self, cross_validation: typing.Optional[skl_model_sel.BaseCrossValidator] = None): """ Prepares and returns a list of callbacks for early stopping in Optuna optimization. - If early stopping with Optuna is enabled and PyTorch is available, this method - will return a list containing a TerminatorCallback instance. If PyTorch is not - available, it will log a warning and return None. + If early stopping with Optuna is enabled, PyTorch is available, and the + cross-validation strategy supports early termination, this method will return + a list containing a TerminatorCallback instance. + + It returns None in any of the following cases: + - early stopping with Optuna is disabled + - PyTorch is not available (warning is logged) + - hold-out cross-validation is detected (fewer than 2 splits) Returns: typing.Optional[typing.List[TerminatorCallback]]: A list of TerminatorCallback - instances if early stopping is enabled and PyTorch is available, otherwise None. + instances when early stopping can be used, otherwise None. """ callbacks: typing.Optional[typing.List[TerminatorCallback]] = None if self.early_stopping_optuna: @@ -179,7 +188,12 @@ def get_callbacks(self): pip install mother[torch] or uv add mother[torch]""" ) else: - callbacks = [TerminatorCallback()] + if cross_validation is not None and cross_validation.get_n_splits() < 2: + module_logger.warning( + "Optuna early termination requires at least 2 CV splits; disabling callback for hold-out setup" + ) + return None + callbacks = [TerminatorCallback(terminator=Terminator(min_n_trials=40))] return callbacks @handle_metadata_routing @@ -254,7 +268,8 @@ def objective(trial: optuna.trial.Trial) -> float: gc.collect() module_logger.info(f"Trial {trial.number}, cv score: {cv_score}") cv_score_not_na: np.ndarray = cv_score[~np.isnan(cv_score)] - report_cross_validation_scores(trial, list(cv_score_not_na)) + if len(cv_score_not_na) > 1: + report_cross_validation_scores(trial, list(cv_score_not_na)) mean_cv_score: float = cv_score_not_na.mean() return mean_cv_score @@ -284,7 +299,7 @@ def objective(trial: optuna.trial.Trial) -> float: objective, n_trials=self.n_trials_optuna, gc_after_trial=True, - callbacks=self.get_callbacks(), + callbacks=self.get_callbacks(cross_validation=cross_validation), ) if default_parameters != {}: diff --git a/test/unit/test_model_tuner.py b/test/unit/test_model_tuner.py index bb71e29..5988022 100644 --- a/test/unit/test_model_tuner.py +++ b/test/unit/test_model_tuner.py @@ -7,7 +7,7 @@ from sklearn.datasets import make_blobs, make_classification from sklearn.linear_model import Lasso from sklearn.metrics import accuracy_score, make_scorer, mean_squared_error -from sklearn.model_selection import GroupKFold +from sklearn.model_selection import GroupKFold, KFold, PredefinedSplit from sklearn.pipeline import Pipeline from sklearn.preprocessing import MinMaxScaler @@ -18,6 +18,13 @@ ) from mother.optimization.core import MotherTuner +try: + import torch # noqa: F401 + + _TORCH_AVAILABLE = True +except (ImportError, OSError): + _TORCH_AVAILABLE = False + # define a simple model to run the tests with @pytest.fixture() @@ -267,3 +274,63 @@ def test_tune_mother_feature_selection_model_pipeline( evaluated = tuner.study.trials[0].params enqueued = model_pipeline.default_parameters() all(evaluated[k] == enqueued[k] for k in evaluated.keys() & enqueued.keys()) + + +@pytest.mark.skipif( + not _TORCH_AVAILABLE, + reason="Optuna early-stopping callback tests require the optional torch extra (mother[torch]).", +) +class TestGetCallbacks: + def test_holdout_cv_disables_early_termination_callback(self, caplog): + caplog.set_level("WARNING") + tuner = MotherTuner(scorer="neg_mean_squared_error", early_stopping_optuna=True) + holdout_cv = PredefinedSplit(test_fold=[-1, -1, 0, 0]) + + callbacks = tuner.get_callbacks(cross_validation=holdout_cv) + + assert callbacks is None + assert "requires at least 2 CV splits" in caplog.text + + def test_multi_split_cv_keeps_early_termination_callback(self): + tuner = MotherTuner(scorer="neg_mean_squared_error", early_stopping_optuna=True) + kfold_cv = KFold(n_splits=3, shuffle=True, random_state=42) + + callbacks = tuner.get_callbacks(cross_validation=kfold_cv) + + assert callbacks is not None + assert len(callbacks) == 1 + + def test_no_cv_argument_keeps_existing_behavior(self): + tuner = MotherTuner(scorer="neg_mean_squared_error", early_stopping_optuna=True) + + callbacks = tuner.get_callbacks() + + assert callbacks is not None + assert len(callbacks) == 1 + + def test_optimize_works_with_holdout_cv(self, lasso_pipeline, synthetic_data, caplog): + caplog.set_level("WARNING") + X, y, _ = synthetic_data + test_fold = np.full(len(X), -1) + test_fold[-20:] = 0 + holdout_cv = PredefinedSplit(test_fold=test_fold) + + tuner = MotherTuner( + scorer="neg_mean_squared_error", + early_stopping_optuna=True, + n_trials_optuna=1, + ) + + model = tuner.optimize( + lasso_pipeline, + X=X, + y=y, + hyperparameter_space_function=lasso_pipeline.get_hyperparameter_space, + default_parameters=lasso_pipeline.default_parameters(), + cross_validation=holdout_cv, + ) + + assert model is not None + assert tuner.study is not None + assert len(tuner.study.trials) >= 1 + assert "requires at least 2 CV splits" in caplog.text diff --git a/test/unit/test_tabpfn.py b/test/unit/test_tabpfn.py index d2a5a14..29df275 100644 --- a/test/unit/test_tabpfn.py +++ b/test/unit/test_tabpfn.py @@ -208,6 +208,7 @@ def test_prefitted_model(self): ) result_prefitted = transformer_prefitted.transform(self.X) + assert transformer_prefitted.model.use_autocast_ is False transformer_newfit = TabPFNEmbeddingTransformer( model_type="regression", use_kfold=False, n_estimators=1, random_state=0