autogluon.tabular 1.4.1b20251014__py3-none-any.whl → 1.5.0b20251222__py3-none-any.whl

autogluon/tabular/models/lgb/lgb_model.py

@@ -46,6 +46,8 @@ class LGBModel(AbstractModel):
     ag_priority_by_problem_type = MappingProxyType({
         SOFTCLASS: 100
     })
+    seed_name = "seed"
+    seed_name_alt = ["seed_value", "random_seed", "random_state"]
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -101,10 +103,46 @@ class LGBModel(AbstractModel):
             Scales linearly with the number of estimators, number of classes, and number of leaves.
             Memory usage peaks during model saving, with the peak consuming approximately 2-4x the size of the model in memory.
         """
+        data_mem_usage = get_approximate_df_mem_usage(X).sum()
+        return cls._estimate_memory_usage_common(
+            num_features=X.shape[1],
+            data_mem_usage=data_mem_usage,
+            hyperparameters=hyperparameters,
+            num_classes=num_classes,
+        )
+
+    @classmethod
+    def _estimate_memory_usage_static_lite(
+        cls,
+        num_samples: int,
+        num_features: int,
+        num_bytes_per_cell: float = 4,
+        hyperparameters: dict = None,
+        num_classes: int = 1,
+        **kwargs,
+    ) -> int:
+        data_mem_usage = num_samples * num_features * num_bytes_per_cell
+        return cls._estimate_memory_usage_common(
+            num_features=num_features,
+            data_mem_usage=data_mem_usage,
+            hyperparameters=hyperparameters,
+            num_classes=num_classes,
+        )
+
+    @classmethod
+    def _estimate_memory_usage_common(
+        cls,
+        num_features: int,
+        data_mem_usage: int | float,
+        hyperparameters: dict | None = None,
+        num_classes: int = 1,
+    ) -> int:
+        """
+        Utility method to avoid code duplication
+        """
         if hyperparameters is None:
             hyperparameters = {}
         num_classes = num_classes if num_classes else 1  # num_classes could be None after initialization if it's a regression problem
-        data_mem_usage = get_approximate_df_mem_usage(X).sum()
         data_mem_usage_bytes = data_mem_usage * 5 + data_mem_usage / 4 * num_classes  # TODO: Extremely crude approximation, can be vastly improved
 
         n_trees_per_estimator = num_classes if num_classes > 2 else 1
@@ -112,7 +150,7 @@ class LGBModel(AbstractModel):
         max_bins = hyperparameters.get("max_bins", 255)
         num_leaves = hyperparameters.get("num_leaves", 31)
         # Memory usage of histogram based on https://github.com/microsoft/LightGBM/issues/562#issuecomment-304524592
-        histogram_mem_usage_bytes = 20 * max_bins * len(X.columns) * num_leaves
+        histogram_mem_usage_bytes = 20 * max_bins * num_features * num_leaves
         histogram_mem_usage_bytes_max = hyperparameters.get("histogram_pool_size", None)
         if histogram_mem_usage_bytes_max is not None:
             histogram_mem_usage_bytes_max *= 1e6  # Convert megabytes to bytes, `histogram_pool_size` is in MB.
@@ -122,18 +160,11 @@ class LGBModel(AbstractModel):
 
         mem_size_per_estimator = n_trees_per_estimator * num_leaves * 100  # very rough estimate
         n_estimators = hyperparameters.get("num_boost_round", DEFAULT_NUM_BOOST_ROUND)
-        n_estimators_min = min(n_estimators, 1000)
-        mem_size_estimators = n_estimators_min * mem_size_per_estimator  # memory estimate after fitting up to 1000 estimators
+        n_estimators_min = min(n_estimators, 5000)
+        mem_size_estimators = n_estimators_min * mem_size_per_estimator  # memory estimate after fitting up to 5000 estimators
 
         approx_mem_size_req = data_mem_usage_bytes + histogram_mem_usage_bytes + mem_size_estimators
-        return approx_mem_size_req
-
-    def _get_random_seed_from_hyperparameters(self, hyperparameters: dict) -> int | None | str:
-        if "seed_value" in hyperparameters:
-            return hyperparameters["seed_value"]
-        if "seed" in hyperparameters:
-            return hyperparameters["seed"]
-        return "N/A"
+        return int(approx_mem_size_req)
 
     def _fit(self, X, y, X_val=None, y_val=None, time_limit=None, num_gpus=0, num_cpus=0, sample_weight=None, sample_weight_val=None, verbosity=2, **kwargs):
         try_import_lightgbm()  # raise helpful error message if LightGBM isn't installed
@@ -292,8 +323,6 @@ class LGBModel(AbstractModel):
         elif self.problem_type == QUANTILE:
             train_params["params"]["quantile_levels"] = self.quantile_levels
 
-        train_params["params"]["seed"] = self.random_seed
-
         # Train LightGBM model:
         # Note that self.model contains a <class 'lightgbm.basic.Booster'> not a LightBGMClassifier or LightGBMRegressor object
         from lightgbm.basic import LightGBMError
@@ -378,6 +407,9 @@ class LGBModel(AbstractModel):
         X = self.preprocess(X, **kwargs)
 
         y_pred_proba = self.model.predict(X, num_threads=num_cpus)
+        return self._post_process_predictions(y_pred_proba=y_pred_proba)
+
+    def _post_process_predictions(self, y_pred_proba) -> np.ndarray:
         if self.problem_type == QUANTILE:
             # y_pred_proba is a pd.DataFrame, need to convert
             y_pred_proba = y_pred_proba.to_numpy()
@@ -430,7 +462,7 @@ class LGBModel(AbstractModel):
         self,
         X: DataFrame,
         y: Series,
-        params,
+        params: dict,
         X_val=None,
         y_val=None,
         X_test=None,
@@ -439,11 +471,14 @@ class LGBModel(AbstractModel):
         sample_weight_val=None,
         sample_weight_test=None,
         save=False,
+        init_train=None,
+        init_val=None,
+        init_test=None,
     ):
         lgb_dataset_params_keys = ["two_round"]  # Keys that are specific to lightGBM Dataset object construction.
         data_params = {key: params[key] for key in lgb_dataset_params_keys if key in params}.copy()
 
-        X = self.preprocess(X, is_train=True)
+        X = self.preprocess(X, y=y, is_train=True)
         if X_val is not None:
             X_val = self.preprocess(X_val)
         if X_test is not None:
@@ -465,7 +500,13 @@ class LGBModel(AbstractModel):
 
         # X, W_train = self.convert_to_weight(X=X)
         dataset_train = construct_dataset(
-            x=X, y=y, location=os.path.join("self.path", "datasets", "train"), params=data_params, save=save, weight=sample_weight
+            x=X,
+            y=y,
+            location=os.path.join("self.path", "datasets", "train"),
+            params=data_params,
+            save=save,
+            weight=sample_weight,
+            init_score=init_train,
         )
         # dataset_train = construct_dataset_lowest_memory(X=X, y=y, location=self.path + 'datasets/train', params=data_params)
         if X_val is not None:
@@ -478,6 +519,7 @@ class LGBModel(AbstractModel):
                 params=data_params,
                 save=save,
                 weight=sample_weight_val,
+                init_score=init_val,
             )
             # dataset_val = construct_dataset_lowest_memory(X=X_val, y=y_val, location=self.path + 'datasets/val', reference=dataset_train, params=data_params)
         else:
@@ -492,6 +534,7 @@ class LGBModel(AbstractModel):
                 params=data_params,
                 save=save,
                 weight=sample_weight_test,
+                init_score=init_test,
             )
         else:
             dataset_test = None

autogluon/tabular/models/lgb/lgb_utils.py

@@ -104,11 +104,11 @@ def softclass_lgbobj(preds, train_data):
     return grad.flatten("F"), hess.flatten("F")
 
 
-def construct_dataset(x: DataFrame, y: Series, location=None, reference=None, params=None, save=False, weight=None):
+def construct_dataset(x: DataFrame, y: Series, location=None, reference=None, params=None, save=False, weight=None, init_score=None):
     try_import_lightgbm()
     import lightgbm as lgb
 
-    dataset = lgb.Dataset(data=x, label=y, reference=reference, free_raw_data=True, params=params, weight=weight)
+    dataset = lgb.Dataset(data=x, label=y, reference=reference, free_raw_data=True, params=params, weight=weight, init_score=init_score)
 
     if save:
         assert location is not None

autogluon/tabular/models/lr/lr_model.py

@@ -43,6 +43,7 @@ class LinearModel(AbstractModel):
     ag_key = "LR"
     ag_name = "LinearModel"
     ag_priority = 30
+    seed_name = "random_state"
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -162,9 +163,6 @@ class LinearModel(AbstractModel):
         for param, val in default_params.items():
             self._set_default_param_value(param, val)
 
-    def _get_random_seed_from_hyperparameters(self, hyperparameters: dict) -> int | None | str:
-        return hyperparameters.get("random_seed", "N/A")
-
     def _get_default_searchspace(self):
         return get_default_searchspace(self.problem_type)
 
@@ -218,7 +216,7 @@ class LinearModel(AbstractModel):
         total_iter = 0
         total_iter_used = 0
         total_max_iter = sum(max_iter_list)
-        model = model_cls(max_iter=max_iter_list[0], random_state=self.random_seed, **params)
+        model = model_cls(max_iter=max_iter_list[0], **params)
         early_stop = False
         for i, cur_max_iter in enumerate(max_iter_list):
             if time_left is not None and (i > 0):

autogluon/tabular/models/lr/lr_preprocessing_utils.py

@@ -5,20 +5,19 @@ from autogluon.features.generators import OneHotEncoderFeatureGenerator
 
 class OheFeaturesGenerator(BaseEstimator, TransformerMixin):
     def __init__(self):
-        self._feature_names = []
-        self._encoder = None
+        pass
 
     def fit(self, X, y=None):
-        self._encoder = OneHotEncoderFeatureGenerator(max_levels=10000, verbosity=0)
-        self._encoder.fit(X)
-        self._feature_names = self._encoder.features_out
+        self.encoder_ = OneHotEncoderFeatureGenerator(max_levels=10000, verbosity=0)
+        self.encoder_.fit(X)
+        self.feature_names_ = self.encoder_.features_out
         return self
 
     def transform(self, X, y=None):
-        return self._encoder.transform_ohe(X)
+        return self.encoder_.transform_ohe(X)
 
     def get_feature_names(self):
-        return self._feature_names
+        return self.feature_names_
 
 
 class NlpDataPreprocessor(BaseEstimator, TransformerMixin):

autogluon/tabular/models/mitra/_internal/core/trainer_finetune.py

@@ -73,6 +73,20 @@ class TrainerFinetune(BaseEstimator):
 
         self.metric = self.cfg.hyperparams['metric']
 
+    def set_device(self, device: str):
+        self.device = device
+        self.model = self.model.to(device=device, non_blocking=True)
+
+    def post_fit_optimize(self):
+        # Minimize memory usage post-fit
+        self.checkpoint = None
+        self.optimizer = None
+        self.scaler = None
+        self.scheduler_warmup = None
+        self.scheduler_reduce_on_plateau = None
+        self.loss = None
+        self.early_stopping = None
+        self.metric = None
 
     def train(self, x_train: np.ndarray, y_train: np.ndarray, x_val: np.ndarray, y_val: np.ndarray):
 
@@ -184,7 +198,6 @@ class TrainerFinetune(BaseEstimator):
 
         self.checkpoint.set_to_best(self.model)
 
-
     def evaluate(self, x_support: np.ndarray, y_support: np.ndarray, x_query: np.ndarray, y_query: np.ndarray) -> PredictionMetrics:
 
         self.model.eval()

autogluon/tabular/models/mitra/mitra_model.py

@@ -2,19 +2,21 @@ from __future__ import annotations
 
 import logging
 import os
+from pathlib import Path
 from typing import List, Optional
 
 import pandas as pd
+from typing_extensions import Self
 
 from autogluon.common.utils.resource_utils import ResourceManager
-from autogluon.core.models import AbstractModel
+from autogluon.tabular.models.abstract.abstract_torch_model import AbstractTorchModel
 from autogluon.features.generators import LabelEncoderFeatureGenerator
 from autogluon.tabular import __version__
 
 logger = logging.getLogger(__name__)
 
 
-class MitraModel(AbstractModel):
+class MitraModel(AbstractTorchModel):
     """
     Mitra is a tabular foundation model pre-trained purely on synthetic data with the goal
     of optimizing fine-tuning performance over in-context learning performance.
@@ -32,6 +34,7 @@ class MitraModel(AbstractModel):
     ag_name = "Mitra"
     weights_file_name = "model.pt"
     ag_priority = 55
+    seed_name = "seed"
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -77,9 +80,6 @@ class MitraModel(AbstractModel):
 
         return X
 
-    def _get_random_seed_from_hyperparameters(self, hyperparameters: dict) -> int | None | str:
-        return hyperparameters.get("seed", "N/A")
-
     def _fit(
         self,
         X: pd.DataFrame,
@@ -157,16 +157,13 @@ class MitraModel(AbstractModel):
         if "verbose" not in hyp:
             hyp["verbose"] = verbosity >= 3
 
-        self.model = model_cls(
-            seed=self.random_seed,
-            **hyp,
-        )
+        self.model = model_cls(**hyp)
 
         X = self.preprocess(X, is_train=True)
         if X_val is not None:
             X_val = self.preprocess(X_val)
 
-        self.model = self.model.fit(
+        model = self.model.fit(
             X=X,
             y=y,
             X_val=X_val,
@@ -174,6 +171,11 @@ class MitraModel(AbstractModel):
             time_limit=time_limit,
         )
 
+        for i in range(len(model.trainers)):
+            model.trainers[i].post_fit_optimize()
+
+        self.model = model
+
         if need_to_reset_torch_threads:
             torch.set_num_threads(torch_threads_og)
 
@@ -195,42 +197,63 @@ class MitraModel(AbstractModel):
         )
         return default_auxiliary_params
 
-    @property
-    def weights_path(self) -> str:
-        return os.path.join(self.path, self.weights_file_name)
+    def weights_path(self, path: str | None = None) -> str:
+        if path is None:
+            path = self.path
+        return str(Path(path) / self.weights_file_name)
 
     def save(self, path: str = None, verbose=True) -> str:
         _model_weights_list = None
         if self.model is not None:
+            self._save_model_artifact(path=path)
             _model_weights_list = []
             for i in range(len(self.model.trainers)):
                 _model_weights_list.append(self.model.trainers[i].model)
-                self.model.trainers[i].checkpoint = None
                 self.model.trainers[i].model = None
-                self.model.trainers[i].optimizer = None
-                self.model.trainers[i].scheduler_warmup = None
-                self.model.trainers[i].scheduler_reduce_on_plateau = None
-            self._weights_saved = True
+
         path = super().save(path=path, verbose=verbose)
         if _model_weights_list is not None:
-            import torch
-
-            os.makedirs(self.path, exist_ok=True)
-            torch.save(_model_weights_list, self.weights_path)
             for i in range(len(self.model.trainers)):
                 self.model.trainers[i].model = _model_weights_list[i]
         return path
 
+    def _save_model_artifact(self, path: str | None):
+        if path is None:
+            path = self.path
+        import torch
+        device_og = self.device
+        self.set_device("cpu")
+
+        _model_weights_list = []
+        for i in range(len(self.model.trainers)):
+            _model_weights_list.append(self.model.trainers[i].model)
+
+        os.makedirs(path, exist_ok=True)
+        torch.save(_model_weights_list, self.weights_path(path=path))
+        self.set_device(device_og)
+        self._weights_saved = True
+
+    def _load_model_artifact(self):
+        import torch
+        device = self.suggest_device_infer()
+        model_weights_list = torch.load(self.weights_path(), weights_only=False)  # nosec B614
+        for i in range(len(self.model.trainers)):
+            self.model.trainers[i].model = model_weights_list[i]
+        self.set_device(device)
+
+    def _set_device(self, device: str):
+        for i in range(len(self.model.trainers)):
+            self.model.trainers[i].set_device(device)
+
+    def get_device(self) -> str:
+        return self.model.trainers[0].device
+
     @classmethod
-    def load(cls, path: str, reset_paths=False, verbose=True):
+    def load(cls, path: str, reset_paths=True, verbose=True) -> Self:
         model: MitraModel = super().load(path=path, reset_paths=reset_paths, verbose=verbose)
 
         if model._weights_saved:
-            import torch
-
-            model_weights_list = torch.load(model.weights_path, weights_only=False)  # nosec B614
-            for i in range(len(model.model.trainers)):
-                model.model.trainers[i].model = model_weights_list[i]
+            model._load_model_artifact()
             model._weights_saved = False
         return model
 
@@ -375,9 +398,12 @@ class MitraModel(AbstractModel):
         return int(gpu_memory_mb * 1e6)
 
     @classmethod
-    def _class_tags(cls) -> dict:
+    def _class_tags(cls):
         return {
             "can_estimate_memory_usage_static": True,
+            "can_set_device": True,
+            "set_device_on_save_to": None,
+            "set_device_on_load": False,
         }
 
     def _more_tags(self) -> dict:

autogluon/tabular/models/realmlp/realmlp_model.py

@@ -16,7 +16,7 @@ from sklearn.impute import SimpleImputer
 
 from autogluon.common.utils.pandas_utils import get_approximate_df_mem_usage
 from autogluon.common.utils.resource_utils import ResourceManager
-from autogluon.core.models import AbstractModel
+from autogluon.tabular.models.abstract.abstract_torch_model import AbstractTorchModel
 from autogluon.tabular import __version__
 
 logger = logging.getLogger(__name__)
@@ -34,7 +34,7 @@ def set_logger_level(logger_name: str, level: int):
 
 
 # pip install pytabkit
-class RealMLPModel(AbstractModel):
+class RealMLPModel(AbstractTorchModel):
     """
     RealMLP is an improved multilayer perception (MLP) model
     through a bag of tricks and better default hyperparameters.
@@ -51,6 +51,7 @@ class RealMLPModel(AbstractModel):
     ag_key = "REALMLP"
     ag_name = "RealMLP"
     ag_priority = 75
+    seed_name = "random_state"
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -82,8 +83,11 @@ class RealMLPModel(AbstractModel):
                 model_cls = RealMLP_TD_S_Regressor
         return model_cls
 
-    def _get_random_seed_from_hyperparameters(self, hyperparameters: dict) -> int | None | str:
-        return hyperparameters.get("random_state", "N/A")
+    def get_device(self) -> str:
+        return self.model.device
+
+    def _set_device(self, device: str):
+        self.model.to(device)
 
     def _fit(
         self,
@@ -178,7 +182,6 @@ class RealMLPModel(AbstractModel):
         self.model = model_cls(
             n_threads=num_cpus,
             device=device,
-            random_state=self.random_seed,
             **init_kwargs,
             **hyp,
         )

autogluon/tabular/models/rf/rf_model.py

@@ -30,6 +30,7 @@ class RFModel(AbstractModel):
     ag_key = "RF"
     ag_name = "RandomForest"
     ag_priority = 80
+    seed_name = "random_state"
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -107,9 +108,6 @@ class RFModel(AbstractModel):
         for param, val in default_params.items():
             self._set_default_param_value(param, val)
 
-    def _get_random_seed_from_hyperparameters(self, hyperparameters: dict) -> int | None | str:
-        return hyperparameters.get("random_state", "N/A")
-
     # TODO: Add in documentation that Categorical default is the first index
     # TODO: enable HPO for RF models
     def _get_default_searchspace(self):
@@ -153,13 +151,13 @@ class RFModel(AbstractModel):
             hyperparameters = {}
         n_estimators_final = hyperparameters.get("n_estimators", 300)
         if isinstance(n_estimators_final, int):
-            n_estimators_minimum = min(40, n_estimators_final)
+            n_estimators = n_estimators_final
         else:  # if search space
-            n_estimators_minimum = 40
+            n_estimators = 40
         num_trees_per_estimator = cls._get_num_trees_per_estimator_static(problem_type=problem_type, num_classes=num_classes)
         bytes_per_estimator = num_trees_per_estimator * len(X) / 60000 * 1e6  # Underestimates by 3x on ExtraTrees
-        expected_min_memory_usage = int(bytes_per_estimator * n_estimators_minimum)
-        return expected_min_memory_usage
+        expected_memory_usage = int(bytes_per_estimator * n_estimators)
+        return expected_memory_usage
 
     def _validate_fit_memory_usage(self, mem_error_threshold: float = 0.5, mem_warning_threshold: float = 0.4, mem_size_threshold: int = 1e7, **kwargs):
         return super()._validate_fit_memory_usage(
@@ -208,7 +206,7 @@ class RFModel(AbstractModel):
             # FIXME: This is inefficient but sklearnex doesn't support computing oob_score after training
             params["oob_score"] = True
 
-        model = model_cls(random_state=self.random_seed, **params)
+        model = model_cls(**params)
 
         time_train_start = time.time()
         for i, n_estimators in enumerate(n_estimator_increments):