autogluon.tabular 1.3.2b20250709__py3-none-any.whl → 1.3.2b20250710__py3-none-any.whl

autogluon/tabular/models/tabm/tabm_model.py

@@ -0,0 +1,275 @@
+"""
+Code Adapted from TabArena: https://github.com/autogluon/tabrepo/blob/main/tabrepo/benchmark/models/ag/tabm/tabm_model.py
+Note: This is a custom implementation of TabM based on TabArena. Because the AutoGluon 1.4 release occurred at nearly
+the same time as TabM became available on PyPi, we chose to use TabArena's implementation
+for the AutoGluon 1.4 release as it has already been benchmarked.
+
+Model: TabM
+Paper: TabM: Advancing Tabular Deep Learning with Parameter-Efficient Ensembling
+Authors: Yury Gorishniy, Akim Kotelnikov, Artem Babenko
+Codebase: https://github.com/yandex-research/tabm
+License: Apache-2.0
+
+Partially adapted from pytabkit's TabM implementation.
+"""
+
+from __future__ import annotations
+
+import logging
+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__
+
+logger = logging.getLogger(__name__)
+
+
+class TabMModel(AbstractModel):
+    ag_key = "TABM"
+    ag_name = "TabM"
+    ag_priority = 85
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._imputer = None
+        self._features_to_impute = None
+        self._features_to_keep = None
+        self._indicator_columns = None
+        self._features_bool = None
+        self._bool_to_cat = None
+
+    def _fit(
+        self,
+        X: pd.DataFrame,
+        y: pd.Series,
+        X_val: pd.DataFrame = None,
+        y_val: pd.Series = None,
+        time_limit: float | None = None,
+        num_cpus: int = 1,
+        num_gpus: float = 0,
+        **kwargs,
+    ):
+        start_time = time.time()
+
+        try:
+            # imports various dependencies such as torch
+            from ._tabm_internal import TabMImplementation
+            from torch.cuda import is_available
+        except ImportError as err:
+            logger.log(
+                40,
+                f"\tFailed to import tabm! To use the TabM model, "
+                f"do: `pip install autogluon.tabular[tabm]=={__version__}`.",
+            )
+            raise err
+
+        device = "cpu" if num_gpus == 0 else "cuda"
+        if (device == "cuda") and (not is_available()):
+            # FIXME: warn instead and switch to CPU.
+            raise AssertionError(
+                "Fit specified to use GPU, but CUDA is not available on this machine. "
+                "Please switch to CPU usage instead.",
+            )
+
+        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,
+                y=y,
+                problem_type=self.problem_type,
+                test_size=0.2,
+                random_state=0,
+            )
+
+        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)
+        if X_val is not None:
+            X_val = self.preprocess(X_val)
+
+        self.model = TabMImplementation(
+            n_threads=num_cpus,
+            device=device,
+            problem_type=self.problem_type,
+            early_stopping_metric=self.stopping_metric,
+            **hyp,
+        )
+
+        self.model.fit(
+            X_train=X,
+            y_train=y,
+            X_val=X_val,
+            y_val=y_val,
+            cat_col_names=X.select_dtypes(include="category").columns.tolist(),
+            time_to_fit_in_seconds=time_limit - (time.time() - start_time) if time_limit is not None else None,
+        )
+
+    # FIXME: bool_to_cat is a hack: Maybe move to abstract model?
+    def _preprocess(
+        self,
+        X: pd.DataFrame,
+        is_train: bool = False,
+        bool_to_cat: bool = False,
+        **kwargs,
+    ) -> pd.DataFrame:
+        """Imputes missing values via the mean and adds indicator columns for numerical features.
+        Converts indicator columns to categorical features to avoid them being treated as numerical by RealMLP.
+        """
+        X = super()._preprocess(X, **kwargs)
+
+        if is_train:
+            self._bool_to_cat = bool_to_cat
+            self._features_bool = self._feature_metadata.get_features(required_special_types=["bool"])
+        if self._bool_to_cat and self._features_bool:
+            # FIXME: Use CategoryFeatureGenerator? Or tell the model which is category
+            X = X.copy(deep=True)
+            X[self._features_bool] = X[self._features_bool].astype("category")
+
+        return X
+
+    def _set_default_params(self):
+        default_params = dict(
+            random_state=0,
+        )
+        for param, val in default_params.items():
+            self._set_default_param_value(param, val)
+
+    @classmethod
+    def supported_problem_types(cls) -> list[str] | None:
+        return ["binary", "multiclass", "regression"]
+
+    def _get_default_stopping_metric(self):
+        return self.eval_metric
+
+    def _get_default_resources(self) -> tuple[int, int]:
+        # only_physical_cores=True is faster in training
+        num_cpus = ResourceManager.get_cpu_count(only_physical_cores=True)
+        num_gpus = min(ResourceManager.get_gpu_count_torch(), 1)
+        return num_cpus, num_gpus
+
+    def _estimate_memory_usage(self, X: pd.DataFrame, **kwargs) -> int:
+        hyperparameters = self._get_model_params()
+        return self.estimate_memory_usage_static(
+            X=X,
+            problem_type=self.problem_type,
+            num_classes=self.num_classes,
+            hyperparameters=hyperparameters,
+            **kwargs,
+        )
+
+    @classmethod
+    def _estimate_memory_usage_static(
+        cls,
+        *,
+        X: pd.DataFrame,
+        hyperparameters: dict = None,
+        num_classes: int | None = 1,
+        **kwargs,
+    ) -> int:
+        """
+        Heuristic memory estimate that correlates strongly with RealMLP
+        """
+        if num_classes is None:
+            num_classes = 1
+        if hyperparameters is None:
+            hyperparameters = {}
+
+        cat_sizes = []
+        for col in X.select_dtypes(include=["category", "object"]):
+            if isinstance(X[col], pd.CategoricalDtype):
+                # Use .cat.codes for category dtype
+                unique_codes = X[col].cat.codes.unique()
+            else:
+                # For object dtype, treat unique strings as codes
+                unique_codes = X[col].astype("category").cat.codes.unique()
+            cat_sizes.append(len(unique_codes))
+
+        n_numerical = len(X.select_dtypes(include=["number"]).columns)
+
+        # TODO: This estimates very high memory usage,
+        #  we probably need to adjust batch size automatically to compensate
+        mem_estimate_bytes = cls._estimate_tabm_ram(
+            hyperparameters=hyperparameters,
+            n_numerical=n_numerical,
+            cat_sizes=cat_sizes,
+            n_classes=num_classes,
+            n_samples=len(X),
+        )
+
+        return mem_estimate_bytes
+
+    @classmethod
+    def _estimate_tabm_ram(
+        cls,
+        hyperparameters: dict,
+        n_numerical: int,
+        cat_sizes: list[int],
+        n_classes: int,
+        n_samples: int,
+    ) -> int:
+        num_emb_n_bins = hyperparameters.get("num_emb_n_bins", 48)
+        d_embedding = hyperparameters.get("d_embedding", 16)
+        d_block = hyperparameters.get("d_block", 512)
+        # not completely sure if this is hidden blocks or all blocks, taking the safe option below
+        n_blocks = hyperparameters.get("n_blocks", "auto")
+        if isinstance(n_blocks, str) and n_blocks == "auto":
+            n_blocks = 3
+        batch_size = hyperparameters.get("batch_size", "auto")
+        if isinstance(batch_size, str) and batch_size == "auto":
+            batch_size = cls.get_tabm_auto_batch_size(n_samples=n_samples)
+        tabm_k = hyperparameters.get("tabm_k", 32)
+        predict_batch_size = hyperparameters.get("eval_batch_size", 1024)
+
+        # not completely sure
+        n_params_num_emb = n_numerical * (num_emb_n_bins + 1) * d_embedding
+        n_params_mlp = (n_numerical + sum(cat_sizes)) * d_embedding * (d_block + tabm_k) \
+                       + (n_blocks - 1) * d_block ** 2 \
+                       + n_blocks * d_block + d_block * (1 + max(1, n_classes))
+        # 4 bytes per float, up to 5 copies of parameters (1 standard, 1 .grad, 2 adam, 1 best_epoch)
+        mem_params = 4 * 5 * (n_params_num_emb + n_params_mlp)
+
+        # compute number of floats in forward pass (per batch element)
+        # todo: numerical embedding layer (not sure if this is entirely correct)
+        n_floats_forward = n_numerical * (num_emb_n_bins + d_embedding)
+        # before and after scale
+        n_floats_forward += 2 * (sum(cat_sizes) + n_numerical * d_embedding)
+        # 2 for pre-act, post-act
+        n_floats_forward += n_blocks * 2 * d_block + 2 * max(1, n_classes)
+        # 2 for forward and backward, 4 bytes per float
+        mem_forward_backward = 4 * max(batch_size * 2, predict_batch_size) * n_floats_forward * tabm_k
+        # * 8 is pessimistic for the long tensors in the forward pass, 4 would probably suffice
+
+        mem_ds = n_samples * (4 * n_numerical + 8 * len(cat_sizes))
+
+        # some safety constants and offsets (the 5 is probably excessive)
+        mem_total = 5 * mem_ds + 1.2 * mem_forward_backward + 1.2 * mem_params + 0.3 * (1024 ** 3)
+
+        return mem_total
+
+    @classmethod
+    def get_tabm_auto_batch_size(cls, n_samples: int) -> int:
+        # by Yury Gorishniy, inferred from the choices in the TabM paper.
+        if n_samples < 2_800:
+            return 32
+        if n_samples < 4_500:
+            return 64
+        if n_samples < 6_400:
+            return 128
+        if n_samples < 32_000:
+            return 256
+        if n_samples < 108_000:
+            return 512
+        return 1024
+
+    @classmethod
+    def _class_tags(cls):
+        return {"can_estimate_memory_usage_static": True}
+
+    def _more_tags(self) -> dict:
+        # TODO: Need to add train params support, track best epoch
+        #  How to force stopping at a specific epoch?
+        return {"can_refit_full": False}