update for rebase, add loading of X_data in ensemble builder

Author: mfeurer

autosklearn/ensemble_building/builder.py

@@ -18,6 +18,7 @@
 from sklearn.utils.validation import check_random_state
 
 from autosklearn.automl_common.common.utils.backend import Backend
+from autosklearn.data.validation import SUPPORTED_FEAT_TYPES
 from autosklearn.data.xy_data_manager import XYDataManager
 from autosklearn.ensemble_building.run import Run, RunID
 from autosklearn.ensembles.abstract_ensemble import AbstractEnsemble
@@ -169,7 +170,9 @@ def __init__(
 
         # Data we may need
         datamanager: XYDataManager = self.backend.load_datamanager()
+        self._X_test: SUPPORTED_FEAT_TYPES | None = datamanager.data.get("X_test", None)
         self._y_test: np.ndarray | None = datamanager.data.get("Y_test", None)
+        self._X_ensemble: SUPPORTED_FEAT_TYPES | None = None
         self._y_ensemble: np.ndarray | None = None
 
     @property
@@ -226,6 +229,29 @@ def targets(self, kind: str = "ensemble") -> np.ndarray | None:
         else:
             raise NotImplementedError(kind)
 
+    def X_data(self, kind: str = "ensemble") -> SUPPORTED_FEAT_TYPES:
+        """The ensemble targets used for training the ensemble
+
+        It will attempt to load and cache them in memory but
+        return None if it can't.
+
+        Returns
+        -------
+        np.ndarray | None
+            The ensemble targets, if they can be loaded
+        """
+        if kind == "ensemble":
+            if self._X_ensemble is None:
+                if os.path.exists(self.backend._get_input_ensemble_filename()):
+                    self._X_ensemble = self.backend.load_input_ensemble()
+            return self._X_ensemble
+
+        elif kind == "test":
+            return self._X_test
+
+        else:
+            raise NotImplementedError(kind)
+
     def run(
         self,
         iteration: int,
@@ -424,7 +450,10 @@ def main(
         for run in requires_update:
             run.record_modified_times()  # So we don't count as modified next time
             run.losses = {
-                metric.name: self.loss(run, metric=metric) for metric in self.metrics
+                metric.name: self.loss(
+                    run, metric=metric, X_data=self.X_data("ensemble")
+                )
+                for metric in self.metrics
             }
 
         # Get the dummy and real runs
@@ -520,9 +549,11 @@ def main(
             return self.ensemble_history, self.ensemble_nbest
 
         targets = cast(np.ndarray, self.targets("ensemble"))  # Sure they exist
+        X_data = self.X_data("ensemble")
 
         ensemble = self.fit_ensemble(
-            candidates,
+            candidates=candidates,
+            X_data=X_data,
             targets=targets,
             runs=runs,
             ensemble_class=self.ensemble_class,
@@ -556,12 +587,14 @@ def main(
 
             run_preds = [r.predictions(kind, precision=self.precision) for r in models]
             pred = ensemble.predict(run_preds)
+            X_data = self.X_data(kind)
 
             scores = calculate_scores(
                 solution=pred_targets,
                 prediction=pred,
                 task_type=self.task_type,
                 metrics=self.metrics,
+                X_data=X_data,
                 scoring_functions=None,
             )
             performance_stamp[f"ensemble_{score_name}_score"] = scores[
@@ -773,6 +806,7 @@ def candidate_selection(
     def fit_ensemble(
         self,
         candidates: list[Run],
+        X_data: SUPPORTED_FEAT_TYPES,
         targets: np.ndarray,
         *,
         runs: list[Run],
@@ -794,6 +828,9 @@ def fit_ensemble(
         candidates: list[Run]
             List of runs to build an ensemble from
 
+        X_data: SUPPORTED_FEAT_TYPES
+            The base level data.
+
         targets: np.ndarray
             The targets to build the ensemble with
 
@@ -851,6 +888,7 @@ def fit_ensemble(
 
         ensemble.fit(
             base_models_predictions=predictions_train,
+            X_data=X_data,
             true_targets=targets,
             model_identifiers=[run.id for run in candidates],
             runs=runs,
@@ -953,7 +991,13 @@ def requires_deletion(
 
         return keep, delete
 
-    def loss(self, run: Run, metric: Scorer, kind: str = "ensemble") -> float:
+    def loss(
+        self,
+        run: Run,
+        metric: Scorer,
+        X_data: SUPPORTED_FEAT_TYPES,
+        kind: str = "ensemble",
+    ) -> float:
         """Calculate the loss for a run
 
         Parameters
@@ -984,6 +1028,7 @@ def loss(self, run: Run, metric: Scorer, kind: str = "ensemble") -> float:
                 prediction=predictions,
                 task_type=self.task_type,
                 metric=metric,
+                X_data=X_data,
             )
         except Exception as e:
             tb = traceback.format_exc()

autosklearn/ensembles/abstract_ensemble.py

@@ -6,6 +6,7 @@
 import numpy as np
 
 from autosklearn.automl_common.common.utils.backend import Backend
+from autosklearn.data.validation import SUPPORTED_FEAT_TYPES
 from autosklearn.ensemble_building.run import Run
 from autosklearn.metrics import Scorer
 from autosklearn.pipeline.base import BasePipeline
@@ -26,6 +27,7 @@ def __init__(
     def fit(
         self,
         base_models_predictions: np.ndarray | List[np.ndarray],
+        X_data: SUPPORTED_FEAT_TYPES,
         true_targets: np.ndarray,
         model_identifiers: List[Tuple[int, int, float]],
         runs: Sequence[Run],
@@ -45,6 +47,8 @@ def fit(
             Can be a list of 2d numpy arrays as well to prevent copying all
             predictions into a single, large numpy array.
 
+        X_data : list-like or sparse data
+
         true_targets : array of shape [n_targets]
 
         model_identifiers : identifier for each base model.

autosklearn/ensembles/ensemble_selection.py

@@ -11,6 +11,7 @@
 
 from autosklearn.automl_common.common.utils.backend import Backend
 from autosklearn.constants import TASK_TYPES
+from autosklearn.data.validation import SUPPORTED_FEAT_TYPES
 from autosklearn.ensemble_building.run import Run
 from autosklearn.ensembles.abstract_ensemble import AbstractEnsemble
 from autosklearn.metrics import Scorer, calculate_losses
@@ -104,6 +105,7 @@ def __getstate__(self) -> Dict[str, Any]:
     def fit(
         self,
         base_models_predictions: List[np.ndarray],
+        X_data: SUPPORTED_FEAT_TYPES,
         true_targets: np.ndarray,
         model_identifiers: List[Tuple[int, int, float]],
         runs: Sequence[Run],
@@ -127,25 +129,31 @@ def fit(
         if self.bagging:
             self._bagging(base_models_predictions, true_targets)
         else:
-            self._fit(base_models_predictions, true_targets)
+            self._fit(
+                predictions=base_models_predictions,
+                X_data=X_data,
+                labels=true_targets,
+            )
         self._calculate_weights()
         self.identifiers_ = model_identifiers
         return self
 
     def _fit(
         self,
         predictions: List[np.ndarray],
+        X_data: SUPPORTED_FEAT_TYPES,
         labels: np.ndarray,
     ) -> EnsembleSelection:
         if self.mode == "fast":
-            self._fast(predictions, labels)
+            self._fast(predictions, X_data, labels)
         else:
-            self._slow(predictions, labels)
+            self._slow(predictions, X_data, labels)
         return self
 
     def _fast(
         self,
         predictions: List[np.ndarray],
+        X_data: SUPPORTED_FEAT_TYPES,
         labels: np.ndarray,
     ) -> None:
         """Fast version of Rich Caruana's ensemble selection method."""
@@ -200,6 +208,7 @@ def _fast(
                     prediction=fant_ensemble_prediction,
                     task_type=self.task_type,
                     metrics=[self.metric],
+                    X_data=X_data,
                     scoring_functions=None,
                 )[self.metric.name]
 
@@ -219,7 +228,12 @@ def _fast(
         self.trajectory_ = trajectory
         self.train_loss_ = trajectory[-1]
 
-    def _slow(self, predictions: List[np.ndarray], labels: np.ndarray) -> None:
+    def _slow(
+        self,
+        predictions: List[np.ndarray],
+        X_data: SUPPORTED_FEAT_TYPES,
+        labels: np.ndarray,
+    ) -> None:
         """Rich Caruana's ensemble selection method."""
         self.num_input_models_ = len(predictions)
 
@@ -242,6 +256,7 @@ def _slow(self, predictions: List[np.ndarray], labels: np.ndarray) -> None:
                     prediction=ensemble_prediction,
                     task_type=self.task_type,
                     metrics=[self.metric],
+                    X_data=X_data,
                     scoring_functions=None,
                 )[self.metric.name]
                 ensemble.pop()

autosklearn/ensembles/singlebest_ensemble.py

@@ -8,6 +8,7 @@
 from smac.runhistory.runhistory import RunHistory
 
 from autosklearn.automl_common.common.utils.backend import Backend
+from autosklearn.data.validation import SUPPORTED_FEAT_TYPES
 from autosklearn.ensemble_building.run import Run
 from autosklearn.ensembles.abstract_ensemble import AbstractEnsemble
 from autosklearn.metrics import Scorer
@@ -52,6 +53,7 @@ def __init__(
     def fit(
         self,
         base_models_predictions: np.ndarray | List[np.ndarray],
+        X_data: SUPPORTED_FEAT_TYPES,
         true_targets: np.ndarray,
         model_identifiers: List[Tuple[int, int, float]],
         runs: Sequence[Run],

autosklearn/metrics/__init__.py

@@ -48,7 +48,7 @@ def __call__(
         self,
         y_true: np.ndarray,
         y_pred: np.ndarray,
-        X_data: Optional[np.ndarray] = None,
+        X_data: Optional[SUPPORTED_XDATA_TYPES] = None,
         sample_weight: Optional[List[float]] = None,
     ) -> float:
         pass
@@ -62,7 +62,7 @@ def __call__(
         self,
         y_true: np.ndarray,
         y_pred: np.ndarray,
-        X_data: Optional[np.ndarray] = None,
+        X_data: Optional[SUPPORTED_XDATA_TYPES] = None,
         sample_weight: Optional[List[float]] = None,
     ) -> float:
         """Evaluate predicted target values for X relative to y_true.
@@ -129,7 +129,7 @@ def __call__(
         self,
         y_true: np.ndarray,
         y_pred: np.ndarray,
-        X_data: Optional[np.ndarray] = None,
+        X_data: Optional[SUPPORTED_XDATA_TYPES] = None,
         sample_weight: Optional[List[float]] = None,
     ) -> float:
         """Evaluate predicted probabilities for X relative to y_true.
@@ -189,7 +189,7 @@ def __call__(
         self,
         y_true: np.ndarray,
         y_pred: np.ndarray,
-        X_data: Optional[np.ndarray] = None,
+        X_data: Optional[SUPPORTED_XDATA_TYPES] = None,
         sample_weight: Optional[List[float]] = None,
     ) -> float:
         """Evaluate decision function output for X relative to y_true.
@@ -563,6 +563,7 @@ def calculate_loss(
     prediction: np.ndarray,
     task_type: int,
     metric: Scorer,
+    X_data: Optional[SUPPORTED_XDATA_TYPES] = None,
 ) -> float:
     """Calculate the loss with a given metric
 
@@ -579,12 +580,16 @@ def calculate_loss(
 
     metric: Scorer
         The metric to use
+
+    X_data: Optional[SUPPORTED_XDATA_TYPES]
+        X data used to obtain the predictions
     """
     losses = calculate_losses(
         solution=solution,
         prediction=prediction,
         task_type=task_type,
         metrics=[metric],
+        X_data=X_data,
     )
     return losses[metric.name]
 
@@ -615,7 +620,7 @@ def calculate_losses(
     metrics: Sequence[Scorer]
         A list of objects that hosts a function to calculate how good the
         prediction is according to the solution.
-    X_data: Optional[np.ndarray]
+    X_data: Optional[SUPPORTED_XDATA_TYPES]
         X data used to obtain the predictions
     scoring_functions: List[Scorer]
         A list of metrics to calculate multiple losses
@@ -652,7 +657,7 @@ def compute_single_metric(
     prediction: np.ndarray,
     solution: np.ndarray,
     task_type: int,
-    X_data: Optional[np.ndarray] = None,
+    X_data: Optional[SUPPORTED_XDATA_TYPES] = None,
 ) -> float:
     """
     Returns a metric for the given Auto-Sklearn Scorer object.

test/test_ensemble_builder/test_ensemble_builder.py

@@ -607,9 +607,10 @@ def test_loss_with_no_ensemble_targets(
     * Should give a loss of np.inf if run has no predictions of a given kind
     """
     run = make_run(predictions=None)
+    X_data = builder.X_data()
     metric = builder.metrics[0]
 
-    assert builder.loss(run, metric=metric, kind=kind) == np.inf
+    assert builder.loss(run, metric=metric, X_data=X_data, kind=kind) == np.inf
 
 
 @parametrize("kind", ["ensemble", "test"])
@@ -623,12 +624,13 @@ def test_loss_with_targets(
     -------
     * Should give a loss < np.inf if the predictions exist
     """
+    X_data = builder.X_data(kind)
     targets = builder.targets(kind)
     metric = builder.metrics[0]
 
     run = make_run(predictions={kind: targets})
 
-    assert builder.loss(run, metric=metric, kind=kind) < np.inf
+    assert builder.loss(run, metric=metric, X_data=X_data, kind=kind) < np.inf
 
 
 def test_delete_runs(builder: EnsembleBuilder, make_run: Callable[..., Run]) -> None:
@@ -680,13 +682,16 @@ def test_fit_ensemble_produces_ensemble(
     -------
     * Should produce an ensemble if all runs have predictions
     """
+    X_data = builder.X_data("ensemble")
     targets = builder.targets("ensemble")
     assert targets is not None
 
     predictions = targets
     runs = [make_run(predictions={"ensemble": predictions}) for _ in range(10)]
 
-    ensemble = builder.fit_ensemble(candidates=runs, targets=targets, runs=runs)
+    ensemble = builder.fit_ensemble(
+        candidates=runs, X_data=X_data, targets=targets, runs=runs
+    )
 
     assert ensemble is not None
 
@@ -701,6 +706,7 @@ def test_fit_with_error_gives_no_ensemble(
     * A run without predictions will raise an error will cause `fit_ensemble` to fail
       as it requires all runs to have valid predictions
     """
+    X_data = builder.X_data("ensemble")
     targets = builder.targets("ensemble")
     assert targets is not None
 
@@ -712,7 +718,7 @@ def test_fit_with_error_gives_no_ensemble(
     runs.append(bad_run)
 
     with pytest.raises(FileNotFoundError):
-        builder.fit_ensemble(candidates=runs, targets=targets, runs=runs)
+        builder.fit_ensemble(candidates=runs, X_data=X_data, targets=targets, runs=runs)
 
 
 @parametrize("time_buffer", [1, 5])