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

autogluon/tabular/models/tabdpt/tabdpt_model.py

@@ -0,0 +1,253 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from autogluon.common.utils.resource_utils import ResourceManager
+from autogluon.core.constants import BINARY, MULTICLASS, REGRESSION
+from autogluon.features.generators import LabelEncoderFeatureGenerator
+from autogluon.tabular.models.abstract.abstract_torch_model import AbstractTorchModel
+
+if TYPE_CHECKING:
+    import numpy as np
+    import pandas as pd
+
+
+# FIXME: Nick:
+#  TODO: batch_size is linear to memory usage
+#   512 default
+#   should be less for very large datasets
+#   128 batch_size on Bioresponse -> 12 GB VRAM
+#       Train Data Rows:    2500
+#       Train Data Columns: 1776
+#       Problem Type:       binary
+#  FIXME: Just set context_size = infinity, everything is way faster, memory usage is way lower, etc.
+#   Train Data Rows:    100000
+#   Train Data Columns: 10
+#   binary
+#   only takes 6.7 GB during inference with batch_size = 512
+# FIXME: Make it work when loading on CPU?
+# FIXME: Can we run 8 in parallel to speed up?
+# TODO: clip_sigma == 1 is terrible, clip_sigma == 16 maybe very good? What about higher values?
+#  clip_sigma >= 16 is roughly all equivalent
+# FIXME: TabDPT stores self.X_test for no reason
+# FIXME: TabDPT creates faiss_knn even if it is never used. Better if `context_size=None` means it is never created.
+# TODO: unit test
+# TODO: memory estimate
+class TabDPTModel(AbstractTorchModel):
+    ag_key = "TABDPT"
+    ag_name = "TabDPT"
+    seed_name = "seed"
+    ag_priority = 50
+    default_random_seed = 0
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._feature_generator = None
+        self._predict_hps = None
+        self._use_flash_og = None
+
+    def _fit(
+        self,
+        X: pd.DataFrame,
+        y: pd.Series,
+        num_cpus: int = 1,
+        num_gpus: int = 0,
+        **kwargs,
+    ):
+        from torch.cuda import is_available
+
+        device = "cuda" if num_gpus != 0 else "cpu"
+        if (device == "cuda") and (not is_available()):
+            raise AssertionError(
+                "Fit specified to use GPU, but CUDA is not available on this machine. "
+                "Please switch to CPU usage instead.",
+            )
+        from tabdpt import TabDPTClassifier, TabDPTRegressor
+
+        model_cls = (
+            TabDPTClassifier
+            if self.problem_type in [BINARY, MULTICLASS]
+            else TabDPTRegressor
+        )
+        fit_params, self._predict_hps = self._get_tabdpt_params(num_gpus=num_gpus)
+
+        X = self.preprocess(X)
+        y = y.to_numpy()
+        self.model = model_cls(
+            device=device,
+            **fit_params,
+        )
+        self.model.fit(X=X, y=y)
+
+    def _get_tabdpt_params(self, num_gpus: float) -> tuple[dict, dict]:
+        model_params = self._get_model_params()
+
+        valid_predict_params = (self.seed_name, "context_size", "permute_classes", "temperature", "n_ensembles")
+
+        predict_params = {}
+        for hp in valid_predict_params:
+            if hp in model_params:
+                predict_params[hp] = model_params.pop(hp)
+        predict_params.setdefault(self.seed_name, self.default_random_seed)
+        predict_params.setdefault("context_size", None)
+
+        supported_predict_params = (
+            (self.seed_name, "context_size", "n_ensembles", "permute_classes", "temperature")
+            if self.problem_type in [BINARY, MULTICLASS]
+            else (self.seed_name, "context_size", "n_ensembles")
+        )
+        predict_params = {key: val for key, val in predict_params.items() if key in supported_predict_params}
+
+        fit_params = model_params
+
+        fit_params.setdefault("verbose", False)
+        fit_params.setdefault("compile", False)
+        if fit_params.get("use_flash", True):
+            fit_params["use_flash"] = self._use_flash(num_gpus=num_gpus)
+        return fit_params, predict_params
+
+    @staticmethod
+    def _use_flash(num_gpus: float) -> bool:
+        """Detect if torch's native flash attention is available on the current machine."""
+        if num_gpus == 0:
+            return False
+
+        import torch
+
+        if not torch.cuda.is_available():
+            return False
+
+        device = torch.device("cuda:0")
+        capability = torch.cuda.get_device_capability(device)
+
+        return capability != (7, 5)
+
+    def _post_fit(self, **kwargs):
+        super()._post_fit(**kwargs)
+        self._use_flash_og = self.model.use_flash
+        return self
+
+    def get_device(self) -> str:
+        return self.model.device
+
+    def _set_device(self, device: str):
+        self.model.to(device)
+        if device == "cpu":
+            self.model.use_flash = False
+            self.model.model.use_flash = False
+        else:
+            self.model.use_flash = self._use_flash_og
+            self.model.model.use_flash = self._use_flash_og
+
+    def _get_default_resources(self) -> tuple[int, int]:
+        # Use only physical cores for better performance based on benchmarks
+        num_cpus = ResourceManager.get_cpu_count(only_physical_cores=True)
+
+        num_gpus = min(1, ResourceManager.get_gpu_count_torch(cuda_only=True))
+
+        return num_cpus, num_gpus
+
+    def get_minimum_resources(
+        self, is_gpu_available: bool = False
+    ) -> dict[str, int | float]:
+        return {
+            "num_cpus": 1,
+            "num_gpus": 0.5 if is_gpu_available else 0,
+        }
+
+    def _predict_proba(self, X, **kwargs) -> np.ndarray:
+        X = self.preprocess(X, **kwargs)
+
+        if self.problem_type in [REGRESSION]:
+            y_pred = self.model.predict(X, **self._predict_hps)
+            return y_pred
+
+        y_pred_proba = self.model.ensemble_predict_proba(X, **self._predict_hps)
+        return self._convert_proba_to_unified_form(y_pred_proba)
+
+    def _preprocess(self, X: pd.DataFrame, **kwargs) -> pd.DataFrame:
+        """TabDPT requires numpy array as input."""
+        X = super()._preprocess(X, **kwargs)
+        if self._feature_generator is None:
+            self._feature_generator = LabelEncoderFeatureGenerator(verbosity=0)
+            self._feature_generator.fit(X=X)
+        if self._feature_generator.features_in:
+            X = X.copy()
+            X[self._feature_generator.features_in] = self._feature_generator.transform(
+                X=X
+            )
+        return X.to_numpy()
+
+    @classmethod
+    def supported_problem_types(cls) -> list[str] | None:
+        return ["binary", "multiclass", "regression"]
+
+    @classmethod
+    def _class_tags(cls):
+        return {"can_estimate_memory_usage_static": True}
+
+    def _more_tags(self) -> dict:
+        return {"can_refit_full": True}
+
+    def _get_default_auxiliary_params(self) -> dict:
+        default_auxiliary_params = super()._get_default_auxiliary_params()
+        default_auxiliary_params.update(
+            {
+                "max_rows": 100000,  # TODO: Test >100k rows
+                "max_features": 2500,  # TODO: Test >2500 features
+                "max_classes": 10,  # TODO: Test >10 classes
+            }
+        )
+        return default_auxiliary_params
+
+    @classmethod
+    def _get_default_ag_args_ensemble(cls, **kwargs) -> dict:
+        default_ag_args_ensemble = super()._get_default_ag_args_ensemble(**kwargs)
+        extra_ag_args_ensemble = {
+            "refit_folds": True,
+        }
+        default_ag_args_ensemble.update(extra_ag_args_ensemble)
+        return default_ag_args_ensemble
+
+    # FIXME: This is copied from TabPFN, but TabDPT is not the same
+    @classmethod
+    def _estimate_memory_usage_static(
+        cls,
+        *,
+        X: pd.DataFrame,
+        hyperparameters: dict | None = None,
+        **kwargs,
+    ) -> int:
+        """Heuristic memory estimate based on TabPFN's memory estimate logic in:
+        https://github.com/PriorLabs/TabPFN/blob/57a2efd3ebdb3886245e4d097cefa73a5261a969/src/tabpfn/model/memory.py#L147.
+
+        This is based on GPU memory usage, but hopefully with overheads it also approximates CPU memory usage.
+        """
+        # TODO: update, this is not correct anymore, consider using internal TabPFN functions directly.
+        features_per_group = 3  # Based on TabPFNv2 default (unused)
+        n_layers = 12  # Based on TabPFNv2 default
+        embedding_size = 192  # Based on TabPFNv2 default
+        dtype_byte_size = 2  # Based on TabPFNv2 default
+
+        model_mem = 14489108  # Based on TabPFNv2 default
+
+        n_samples, n_features = X.shape[0], min(X.shape[1], 500)
+        n_feature_groups = (
+            n_features
+        ) / features_per_group + 1  # TODO: Unsure how to calculate this
+
+        X_mem = n_samples * n_feature_groups * dtype_byte_size
+        activation_mem = (
+            n_samples * n_feature_groups * embedding_size * n_layers * dtype_byte_size
+        )
+
+        baseline_overhead_mem_est = 1e9  # 1 GB generic overhead
+
+        # Add some buffer to each term + 1 GB overhead to be safe
+        memory_estimate = model_mem + 4 * X_mem + 2 * activation_mem + baseline_overhead_mem_est
+
+        # TabDPT memory estimation is very inaccurate because it is using TabPFN memory estimate. Double it to be safe.
+        memory_estimate = memory_estimate * 2
+
+        # Note: This memory estimate is way off if `context_size` is not None
+        return int(memory_estimate)

autogluon/tabular/models/tabicl/tabicl_model.py

@@ -10,14 +10,14 @@ import pandas as pd
 
 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 import __version__
+from autogluon.tabular.models.abstract.abstract_torch_model import AbstractTorchModel
 
 logger = logging.getLogger(__name__)
 
 
 # TODO: Verify if crashes when weights are not yet downloaded and fit in parallel
-class TabICLModel(AbstractModel):
+class TabICLModel(AbstractTorchModel):
     """
     TabICL is a foundation model for tabular data using in-context learning
     that is scalable to larger datasets than TabPFNv2. It is pretrained purely on synthetic data.
@@ -35,6 +35,7 @@ class TabICLModel(AbstractModel):
     ag_key = "TABICL"
     ag_name = "TabICL"
     ag_priority = 65
+    seed_name = "random_state"
 
     def get_model_cls(self):
         from tabicl import TabICLClassifier
@@ -89,7 +90,6 @@ class TabICLModel(AbstractModel):
             **hyp,
             device=device,
             n_jobs=num_cpus,
-            random_state=self.random_seed,
         )
         X = self.preprocess(X)
         self.model = self.model.fit(
@@ -97,8 +97,18 @@ class TabICLModel(AbstractModel):
             y=y,
         )
 
-    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_.type
+
+    # TODO: Better to have an official TabICL method for this
+    def _set_device(self, device: str):
+        device = self.to_torch_device(device)
+        self.model.device_ = device
+        self.model.device = self.model.device_.type
+        self.model.model_ = self.model.model_.to(self.model.device_)
+        self.model.inference_config_.COL_CONFIG.device = self.model.device_
+        self.model.inference_config_.ROW_CONFIG.device = self.model.device_
+        self.model.inference_config_.ICL_CONFIG.device = self.model.device_
 
     def _get_default_auxiliary_params(self) -> dict:
         default_auxiliary_params = super()._get_default_auxiliary_params()

autogluon/tabular/models/tabm/tabm_model.py

@@ -15,13 +15,13 @@ import time
 import pandas as pd
 
 from autogluon.common.utils.resource_utils import ResourceManager
-from autogluon.core.models import AbstractModel
 from autogluon.tabular import __version__
+from autogluon.tabular.models.abstract.abstract_torch_model import AbstractTorchModel
 
 logger = logging.getLogger(__name__)
 
 
-class TabMModel(AbstractModel):
+class TabMModel(AbstractTorchModel):
     """
     TabM is an efficient ensemble of MLPs that is trained simultaneously with mostly shared parameters.
 
@@ -39,6 +39,7 @@ class TabMModel(AbstractModel):
     ag_key = "TABM"
     ag_name = "TabM"
     ag_priority = 85
+    seed_name = "random_state"
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -86,7 +87,7 @@ class TabMModel(AbstractModel):
         if X_val is None:
             from autogluon.core.utils import generate_train_test_split
 
-            X_train, X_val, y_train, y_val = generate_train_test_split(
+            X, X_val, y, y_val = generate_train_test_split(
                 X=X,
                 y=y,
                 problem_type=self.problem_type,
@@ -97,7 +98,7 @@ class TabMModel(AbstractModel):
         hyp = self._get_model_params()
         bool_to_cat = hyp.pop("bool_to_cat", True)
 
-        X = self.preprocess(X, is_train=True, bool_to_cat=bool_to_cat)
+        X = self.preprocess(X, y=y, is_train=True, bool_to_cat=bool_to_cat)
         if X_val is not None:
             X_val = self.preprocess(X_val)
 
@@ -106,7 +107,6 @@ class TabMModel(AbstractModel):
             device=device,
             problem_type=self.problem_type,
             early_stopping_metric=self.stopping_metric,
-            random_state=self.random_seed,
             **hyp,
         )
 
@@ -142,8 +142,13 @@ class TabMModel(AbstractModel):
 
         return X
 
-    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_.type
+
+    def _set_device(self, device: str):
+        device = self.to_torch_device(device)
+        self.model.device_ = device
+        self.model.model_ = self.model.model_.to(device)
 
     @classmethod
     def supported_problem_types(cls) -> list[str] | None:
@@ -258,6 +263,15 @@ class TabMModel(AbstractModel):
 
         return mem_total
 
+    def _get_default_auxiliary_params(self) -> dict:
+        default_auxiliary_params = super()._get_default_auxiliary_params()
+        default_auxiliary_params.update(
+            {
+                "max_batch_size": 16384,  # avoid excessive VRAM usage
+            }
+        )
+        return default_auxiliary_params
+
     @classmethod
     def get_tabm_auto_batch_size(cls, n_samples: int) -> int:
         # by Yury Gorishniy, inferred from the choices in the TabM paper.
@@ -275,7 +289,10 @@ class TabMModel(AbstractModel):
 
     @classmethod
     def _class_tags(cls):
-        return {"can_estimate_memory_usage_static": True}
+        return {
+            "can_estimate_memory_usage_static": True,
+            "reset_torch_threads": True,
+        }
 
     def _more_tags(self) -> dict:
         # TODO: Need to add train params support, track best epoch

autogluon/tabular/models/tabpfnmix/tabpfnmix_model.py

@@ -42,6 +42,7 @@ class TabPFNMixModel(AbstractModel):
     ag_key = "TABPFNMIX"
     ag_name = "TabPFNMix"
     ag_priority = 45
+    seed_name = "random_state"
 
     weights_file_name = "model.pt"
 
@@ -123,6 +124,7 @@ class TabPFNMixModel(AbstractModel):
             raise AssertionError(f"Max allowed classes for the model is {max_classes}, " f"but found {self.num_classes} classes.")
 
         params = self._get_model_params()
+        random_state = params.pop(self.seed_name, self.default_random_seed)
         sample_rows = ag_params.get("sample_rows", None)
         sample_rows_val = ag_params.get("sample_rows_val", None)
         max_rows = ag_params.get("max_rows", None)
@@ -133,11 +135,11 @@ class TabPFNMixModel(AbstractModel):
 
         # TODO: Make sample_rows generic
         if sample_rows is not None and isinstance(sample_rows, int) and len(X) > sample_rows:
-            X, y = self._subsample_data(X=X, y=y, num_rows=sample_rows)
+            X, y = self._subsample_data(X=X, y=y, num_rows=sample_rows, random_state=random_state)
 
         # TODO: Make sample_rows generic
         if X_val is not None and y_val is not None and sample_rows_val is not None and isinstance(sample_rows_val, int) and len(X_val) > sample_rows_val:
-            X_val, y_val = self._subsample_data(X=X_val, y=y_val, num_rows=sample_rows_val)
+            X_val, y_val = self._subsample_data(X=X_val, y=y_val, num_rows=sample_rows_val, random_state=random_state)
 
         from ._internal.core.enums import Task
         if self.problem_type in [REGRESSION, QUANTILE]:
@@ -178,7 +180,7 @@ class TabPFNMixModel(AbstractModel):
         elif weights_path is not None:
             logger.log(15, f'\tLoading pre-trained weights from file... (weights_path="{weights_path}")')
 
-        cfg = ConfigRun(hyperparams=params, task=task, device=device, seed=self.random_seed)
+        cfg = ConfigRun(hyperparams=params, task=task, device=device, seed=random_state)
 
         if cfg.hyperparams["max_epochs"] == 0 and cfg.hyperparams["n_ensembles"] != 1:
             logger.log(
@@ -242,14 +244,14 @@ class TabPFNMixModel(AbstractModel):
         return self
 
     # TODO: Make this generic by creating a generic `preprocess_train` and putting this logic prior to `_preprocess`.
-    def _subsample_data(self, X: pd.DataFrame, y: pd.Series, num_rows: int) -> (pd.DataFrame, pd.Series):
+    def _subsample_data(self, X: pd.DataFrame, y: pd.Series, num_rows: int, random_state: int | None = 0) -> (pd.DataFrame, pd.Series):
         num_rows_to_drop = len(X) - num_rows
         X, _, y, _ = generate_train_test_split(
             X=X,
             y=y,
             problem_type=self.problem_type,
             test_size=num_rows_to_drop,
-            random_state=self.random_seed,
+            random_state=random_state,
             min_cls_count_train=1,
         )
         return X, y