Optimize memory usage, add bag info support

Author: Innixma

tabrepo/benchmark/experiment/experiment_constructor.py

@@ -9,7 +9,7 @@
 from autogluon.core.models import AbstractModel
 
 from tabrepo.benchmark.models.wrapper.abstract_class import AbstractExecModel
-from tabrepo.benchmark.models.wrapper.AutoGluon_class import AGSingleWrapper
+from tabrepo.benchmark.models.wrapper.AutoGluon_class import AGSingleWrapper, AGSingleBagWrapper
 from tabrepo.benchmark.models.wrapper.ag_model import AGModelWrapper
 from tabrepo.benchmark.experiment.experiment_runner import ExperimentRunner, OOFExperimentRunner
 from tabrepo.benchmark.models.model_register import infer_model_cls
@@ -220,6 +220,7 @@ class AGModelExperiment(Experiment):
     experiment_kwargs: dict, optional
         The kwargs passed to the init of `experiment_cls`.
     """
+    _method_cls = AGSingleWrapper
 
     def __init__(
         self,
@@ -251,7 +252,7 @@ def __init__(
         method_kwargs["fit_kwargs"]["raise_on_model_failure"] = raise_on_model_failure
         super().__init__(
             name=name,
-            method_cls=AGSingleWrapper,
+            method_cls=self._method_cls,
             method_kwargs={
                 "model_cls": model_cls,
                 "model_hyperparameters": model_hyperparameters,
@@ -328,6 +329,8 @@ class AGModelBagExperiment(AGModelExperiment):
     method_kwargs: dict, optional
     experiment_kwargs: dict, optional
     """
+    _method_cls = AGSingleBagWrapper
+
     def __init__(
         self,
         name: str,

tabrepo/benchmark/experiment/experiment_runner.py

@@ -3,7 +3,9 @@
 import datetime
 from typing import Literal, Type
 
+import numpy as np
 import pandas as pd
+from pandas.api.types import is_integer_dtype
 
 from autogluon.core.data.label_cleaner import LabelCleaner, LabelCleanerDummy
 from autogluon.core.metrics import get_metric, Scorer
@@ -13,17 +15,18 @@
 from tabrepo.benchmark.models.wrapper.abstract_class import AbstractExecModel
 
 
+# TODO: make a dataclass so type hinter is happy with subclasses?
 class ExperimentRunner:
     def __init__(
         self,
+        *,
         method_cls: Type[AbstractExecModel],
         task: OpenMLTaskWrapper,
         fold: int,
         task_name: str,
         method: str,
         fit_args: dict | None = None,
         cleanup: bool = True,
-        compute_simulation_artifacts: bool = True,
         input_format: Literal["openml", "csv"] = "openml",
         cacher: AbstractCacheFunction | None = None,
     ):
@@ -36,7 +39,6 @@ def __init__(
         self.fit_args = fit_args
         self.cleanup = cleanup
         self.input_format = input_format
-        self.compute_simulation_artifacts = compute_simulation_artifacts
         self.eval_metric_name = ag_eval_metric_map[self.task.problem_type]  # FIXME: Don't hardcode eval metric
         self.eval_metric: Scorer = get_metric(metric=self.eval_metric_name, problem_type=self.task.problem_type)
         self.model = None
@@ -194,6 +196,19 @@ def _cleanup(self):
 
 
 class OOFExperimentRunner(ExperimentRunner):
+    def __init__(
+        self,
+        *,
+        compute_simulation_artifacts: bool = True,
+        compute_bag_info: bool = True,
+        optimize_simulation_artifacts_memory: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.compute_simulation_artifacts = compute_simulation_artifacts
+        self.compute_bag_info = compute_bag_info
+        self.optimize_simulation_artifacts_memory = optimize_simulation_artifacts_memory
+
     def post_evaluate(self, out: dict) -> dict:
         out = super().post_evaluate(out=out)
         if self.compute_simulation_artifacts and self.model.can_get_oof:
@@ -205,16 +220,36 @@ def post_evaluate(self, out: dict) -> dict:
                 if self.task.problem_type == "binary":
                     simulation_artifact["pred_proba_dict_test"] = simulation_artifact["pred_proba_dict_test"].iloc[:, 1]
             simulation_artifact["y_test"] = self.label_cleaner.transform(self.y_test)
+
+            if self.optimize_simulation_artifacts_memory:
+                # optimize memory
+                simulation_artifact["y_test"].index = pd.to_numeric(simulation_artifact["y_test"].index, downcast="integer")
+                simulation_artifact["y_val"].index = pd.to_numeric(simulation_artifact["y_val"].index, downcast="integer")
+
+                simulation_artifact["y_test_idx"] = simulation_artifact["y_test"].index.values
+                simulation_artifact["y_val_idx"] = simulation_artifact["y_val"].index.values
+
+                simulation_artifact["y_test"] = simulation_artifact["y_test"].values
+                simulation_artifact["y_val"] = simulation_artifact["y_val"].values
+                if is_integer_dtype(simulation_artifact["y_test"]):
+                    simulation_artifact["y_test"] = pd.to_numeric(simulation_artifact["y_test"], downcast="integer")
+                if is_integer_dtype(simulation_artifact["y_val"]):
+                    simulation_artifact["y_val"] = pd.to_numeric(simulation_artifact["y_val"], downcast="integer")
+
+                simulation_artifact["pred_proba_dict_test"] = simulation_artifact["pred_proba_dict_test"].astype(np.float32)
+                simulation_artifact["pred_proba_dict_val"] = simulation_artifact["pred_proba_dict_val"].astype(np.float32)
+
+                simulation_artifact["pred_proba_dict_test"] = simulation_artifact["pred_proba_dict_test"].values
+                simulation_artifact["pred_proba_dict_val"] = simulation_artifact["pred_proba_dict_val"].values
+
             simulation_artifact["label"] = self.task.label
             simulation_artifact["metric"] = self.eval_metric_name
 
             out["metric_error_val"] = self.model.get_metric_error_val()
-            # out["metric_error_val"] = evaluate(
-            #     y_true=simulation_artifact["y_val"],
-            #     y_pred=self.label_cleaner.transform(out["predictions"]),
-            #     y_pred_proba=self.label_cleaner.transform_proba(out["probabilities"])
-            # )
-            # out["metric_error_val"] = self.eval_metric.error(simulation_artifact["y_val"], simulation_artifact["pred_proba_dict_val"])
+
+            if self.compute_bag_info and (self.model.can_get_per_child_oof and self.model.can_get_per_child_val_idx):
+                simulation_artifact["bag_info"] = self.model.bag_artifact(X_test=self.X_test)
+
 
             simulation_artifact["pred_proba_dict_val"] = {self.method: simulation_artifact["pred_proba_dict_val"]}
             simulation_artifact["pred_proba_dict_test"] = {self.method: simulation_artifact["pred_proba_dict_test"]}

tabrepo/benchmark/models/wrapper/AutoGluon_class.py

@@ -4,6 +4,7 @@
 import shutil
 from typing import Type
 
+import numpy as np
 import pandas as pd
 
 from tabrepo.benchmark.models.wrapper.abstract_class import AbstractExecModel
@@ -132,6 +133,9 @@ def __init__(
 
         super().__init__(init_kwargs=init_kwargs, fit_kwargs=fit_kwargs, preprocess_data=preprocess_data, preprocess_label=preprocess_label, **kwargs)
 
+    def post_fit(self, X: pd.DataFrame, y: pd.Series, X_test: pd.DataFrame):
+        self.failure_artifact = self.get_metadata_failure()
+
     def get_hyperparameters(self):
         hyperparameters = self.predictor.model_hyperparameters(model=self.predictor.model_best, output_format="user")
         return hyperparameters
@@ -146,22 +150,91 @@ def model_cls(self) -> Type["AbstractModel"]:
             model_cls = ag_model_register.key_to_cls(key=self._model_cls)
         return model_cls
 
-    def get_metadata(self) -> dict:
-        metadata = {}
+    def _load_model(self):
+        model_names = self.predictor.model_names(can_infer=True)
+        assert len(model_names) == 1
+        model_name = self.predictor.model_names()[0]
+        return self.predictor._trainer.load_model(model_name)
 
-        model = self.predictor._trainer.load_model(self.predictor.model_best)
-        metadata["info"] = model.get_info(include_feature_metadata=False)
+    def get_metadata_init(self) -> dict:
+        metadata = {}
         metadata["hyperparameters"] = self.get_hyperparameters()
         metadata["model_cls"] = self.model_cls.__name__
         metadata["model_type"] = self.model_cls.ag_key  # TODO: rename to ag_key?
         metadata["name_prefix"] = self.model_cls.ag_name  # TODO: rename to ag_name?
         metadata["model_hyperparameters"] = self.model_hyperparameters
         metadata["init_kwargs_extra"] = self.init_kwargs_extra
         metadata["fit_kwargs_extra"] = self.fit_kwargs_extra
+        return metadata
+
+    def get_metadata_fit(self) -> dict:
+        metadata = {}
+        model = self.predictor._trainer.load_model(self.predictor.model_best)
+        metadata["info"] = model.get_info(include_feature_metadata=False)
         metadata["disk_usage"] = model.disk_usage()
         metadata["num_cpus"] = model.fit_num_cpus
         metadata["num_gpus"] = model.fit_num_gpus
         metadata["num_cpus_child"] = model.fit_num_cpus_child
         metadata["num_gpus_child"] = model.fit_num_gpus_child
         metadata["fit_metadata"] = model.get_fit_metadata()
         return metadata
+
+    def get_metadata_failure(self) -> dict:
+        metadata = {
+            "model_failures": self.predictor.model_failures()
+        }
+        return metadata
+
+    def get_metadata(self) -> dict:
+        metadata = self.get_metadata_init()
+        metadata_fit = self.get_metadata_fit()
+
+        metadata.update(metadata_fit)
+        return metadata
+
+
+class AGSingleBagWrapper(AGSingleWrapper):
+    can_get_per_child_oof = True
+    can_get_per_child_val_idx = True
+
+    def bag_artifact(self, X_test: pd.DataFrame):
+        model = self._load_model()
+        bag_info = {}
+        bag_info["pred_proba_test_per_child"] = self.get_per_child_test(X_test=X_test, model=model)
+        bag_info["val_idx_per_child"] = self.get_per_child_val_idx(model=model)
+        return bag_info
+
+    def get_per_child_val_idx(self, model=None) -> list[np.ndarray]:
+        if model is None:
+            model = self._load_model()
+        X, y = self.predictor.load_data_internal()
+        all_kfolds = []
+        # TODO: Make this a bagged ensemble method
+        if model._child_oof:
+            all_kfolds = [(None, X.index.values)]
+        else:
+            for n_repeat, k in enumerate(model._k_per_n_repeat):
+                kfolds = model._cv_splitters[n_repeat].split(X=X, y=y)
+                cur_kfolds = kfolds[n_repeat * k: (n_repeat + 1) * k]
+                all_kfolds += cur_kfolds
+
+        val_idx_per_child = []
+        for fold_idx, (train_idx, val_idx) in enumerate(all_kfolds):
+            val_idx = pd.to_numeric(val_idx, downcast="integer")  # memory opt
+            val_idx_per_child.append(val_idx)
+
+        return val_idx_per_child
+
+    # TODO: Can avoid predicting on test twice by doing it all in one go
+    def get_per_child_test(self, X_test: pd.DataFrame, model=None) -> list[np.ndarray]:
+        if model is None:
+            model = self._load_model()
+        X_test_inner = self.predictor.transform_features(data=X_test, model=model.name)
+
+        if model.can_predict_proba():
+            per_child_test_preds = model.predict_proba_children(X=X_test_inner)
+        else:
+            per_child_test_preds = model.predict_children(X=X_test_inner)
+
+        per_child_test_preds = [preds_child.astype(np.float32) for preds_child in per_child_test_preds]  # memory opt
+        return per_child_test_preds

tabrepo/benchmark/models/wrapper/abstract_class.py

@@ -10,6 +10,8 @@
 
 class AbstractExecModel:
     can_get_oof = False
+    can_get_per_child_oof = False
+    can_get_per_child_test = False
 
     # TODO: Prateek: Find a way to put AutoGluon as default - in the case the user does not want their own class
     def __init__(
@@ -25,6 +27,7 @@ def __init__(
         self.preprocess_label = preprocess_label
         self.label_cleaner: LabelCleaner = None
         self._feature_generator = None
+        self.failure_artifact = None
 
     def transform_y(self, y: pd.Series) -> pd.Series:
         return self.label_cleaner.transform(y)
@@ -54,6 +57,9 @@ def _preprocess_fit_transform(self, X: pd.DataFrame, y: pd.Series):
         y = self.transform_y(y)
         return X, y
 
+    def post_fit(self, X: pd.DataFrame, y: pd.Series, X_test: pd.DataFrame):
+        pass
+
     # TODO: Prateek, Add a toggle here to see if user wants to fit or fit and predict, also add model saving functionality
     # TODO: Nick: Temporary name
     def fit_custom(self, X: pd.DataFrame, y: pd.Series, X_test: pd.DataFrame):
@@ -68,6 +74,8 @@ def fit_custom(self, X: pd.DataFrame, y: pd.Series, X_test: pd.DataFrame):
         with (Timer() as timer_fit):
             self.fit(X, y)
 
+        self.post_fit(X=X, y=y, X_test=X_test)
+
         if self.problem_type in ['binary', 'multiclass']:
             with Timer() as timer_predict:
                 y_pred_proba = self.predict_proba(X_test)