autogluon.tabular 1.4.1b20251212__py3-none-any.whl → 1.5.0b20251220__py3-none-any.whl

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.
@@ -161,7 +163,7 @@ class MitraModel(AbstractModel):
         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,
@@ -169,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)
 
@@ -190,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
 
@@ -370,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.
@@ -83,6 +83,12 @@ class RealMLPModel(AbstractModel):
                 model_cls = RealMLP_TD_S_Regressor
         return model_cls
 
+    def get_device(self) -> str:
+        return self.model.device
+
+    def _set_device(self, device: str):
+        self.model.to(device)
+
     def _fit(
         self,
         X: pd.DataFrame,

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.
@@ -97,6 +97,19 @@ class TabICLModel(AbstractModel):
             y=y,
         )
 
+    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()
         default_auxiliary_params.update(

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.
 
@@ -49,7 +49,6 @@ class TabMModel(AbstractModel):
         self._indicator_columns = None
         self._features_bool = None
         self._bool_to_cat = None
-        self.device = None
 
     def _fit(
         self,
@@ -88,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,
@@ -99,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)
 
@@ -143,80 +142,13 @@ class TabMModel(AbstractModel):
 
         return X
 
-    def save(self, path: str = None, verbose=True) -> str:
-        """
-        Need to set device to CPU to be able to load on a non-GPU environment
-        """
-        import torch
-
-        # Save on CPU to ensure the model can be loaded without GPU
-        if self.model is not None:
-            self.device = self.model.device_
-            device_cpu = torch.device("cpu")
-            self.model.model_ = self.model.model_.to(device_cpu)
-            self.model.device_ = device_cpu
-        path = super().save(path=path, verbose=verbose)
-        # Put the model back to the device after the save
-        if self.model is not None:
-            self.model.model_.to(self.device)
-            self.model.device_ = self.device
-
-        return path
-
-    @classmethod
-    def load(cls, path: str, reset_paths=True, verbose=True):
-        """
-        Loads the model from disk to memory.
-        The loaded model will be on the same device it was trained on (cuda/mps);
-        if the device is not available (trained on GPU, deployed on CPU), then `cpu` will be used.
-
-        Parameters
-        ----------
-        path : str
-            Path to the saved model, minus the file name.
-            This should generally be a directory path ending with a '/' character (or appropriate path separator value depending on OS).
-            The model file is typically located in os.path.join(path, cls.model_file_name).
-        reset_paths : bool, default True
-            Whether to reset the self.path value of the loaded model to be equal to path.
-            It is highly recommended to keep this value as True unless accessing the original self.path value is important.
-            If False, the actual valid path and self.path may differ, leading to strange behaviour and potential exceptions if the model needs to load any other files at a later time.
-        verbose : bool, default True
-            Whether to log the location of the loaded file.
-
-        Returns
-        -------
-        model : cls
-            Loaded model object.
-        """
-        import torch
-
-        model: TabMModel = super().load(path=path, reset_paths=reset_paths, verbose=verbose)
-
-        # Put the model on the same device it was trained on (GPU/MPS) if it is available; otherwise use CPU
-        if model.model is not None:
-            original_device_type = model.device.type
-            if "cuda" in original_device_type:
-                # cuda: nvidia GPU
-                device = torch.device(original_device_type if torch.cuda.is_available() else "cpu")
-            elif "mps" in original_device_type:
-                # mps: Apple Silicon
-                device = torch.device(original_device_type if torch.backends.mps.is_available() else "cpu")
-            else:
-                device = torch.device(original_device_type)
-
-            if verbose and (original_device_type != device.type):
-                logger.log(15, f"Model is trained on {original_device_type}, but the device is not available - loading on {device.type}")
-
-            model.set_device(device=device)
+    def get_device(self) -> str:
+        return self.model.device_.type
 
-        return model
-
-    def set_device(self, device):
-        self.device = device
-        if self.model is not None:
-            self.model.device_ = device
-            if self.model.model_ is not None:
-                self.model.model_ = self.model.model_.to(device)
+    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:
@@ -331,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.
@@ -348,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