v1.7.1: move realmlp to different devices, implement missing lgbm parameters, fix xRFM bug

Author: dholzmueller

.gitignore

@@ -11,6 +11,7 @@ public_export
 dist
 files
 plots
+lightning_logs
 
 docs/build
 docs/source/modules.rst

README.md

@@ -28,7 +28,7 @@ on our benchmarks.
     (e.g., `class_error`, `cross_entropy`, `brier`, `1-auc_ovr`), or the corresponding `Regressor`. 
     (This might take very long to fit.)
   - For only a single model, we recommend using 
-    `RealMLP_HPO_Classifier(n_cv=8, hpo_space_name='tabarena', use_caruana_ensembling=True, n_hyperopt_steps=50)`,
+    `RealMLP_HPO_Classifier(n_cv=8, hpo_space_name='tabarena-new', use_caruana_ensembling=True, n_hyperopt_steps=50)`,
     also with `val_metric_name` as above, or the corresponding `Regressor`.
 - **Models**: [TabArena](https://github.com/AutoGluon/tabarena) 
   also includes some newer models like RealMLP and TabM 
@@ -184,9 +184,9 @@ If you use this repository for research purposes, please cite our [paper](https:
 - Léo Grinsztajn (deep learning baselines, plotting)
 - Ingo Steinwart (UCI dataset download)
 - Katharina Strecker (PyTorch-Lightning interface)
+- Daniel Beaglehole (part of the xRFM implementation)
 - Lennart Purucker (some features/fixes)
 - Jérôme Dockès (deployment, continuous integration)
-- 
 
 ## Acknowledgements
 
@@ -200,6 +200,12 @@ and https://docs.ray.io/en/latest/cluster/vms/user-guides/community/slurm.html
 
 ## Releases (see git tags)
 
+- v1.7.1:
+    - LightGBM now processes the `extra_trees`, `max_cat_to_onehot`, and `min_data_per_group` parameters 
+      used in the `'tabarena'` search space, which should improve results.
+    - Scikit-learn interfaces for RealMLP (TD, HPO) now support moving the model to a different device 
+      (e.g., before saving). This can be achived using, e.g., `model.to('cpu')` (which is in-place).
+    - Fixed an xRFM bug in handling binary categorical features.
 - v1.7.0:
     - added [xRFM](https://arxiv.org/abs/2508.10053) (D, HPO)
     - added new `'tabarena-new'` search space for RealMLP-HPO, including per-fold ensembling (more expensive)

pytabkit/__about__.py

@@ -2,4 +2,4 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "1.7.0"
+__version__ = "1.7.1"

pytabkit/bench/alg_wrappers/interface_wrappers.py

@@ -1,4 +1,5 @@
 import shutil
+import time
 from pathlib import Path
 from typing import Callable, List, Optional, Dict
 
@@ -9,6 +10,7 @@
 from pytabkit.models.alg_interfaces.autogluon_model_interfaces import AutoGluonModelAlgInterface
 from pytabkit.models.alg_interfaces.catboost_interfaces import CatBoostSubSplitInterface, CatBoostHyperoptAlgInterface, \
     CatBoostSklearnSubSplitInterface, RandomParamsCatBoostAlgInterface
+from pytabkit.models.alg_interfaces.custom_interfaces import TabPFNV2SubSplitInterface
 from pytabkit.models.alg_interfaces.ensemble_interfaces import PrecomputedPredictionsAlgInterface, \
     CaruanaEnsembleAlgInterface, AlgorithmSelectionAlgInterface
 from pytabkit.models.alg_interfaces.lightgbm_interfaces import LGBMSubSplitInterface, LGBMHyperoptAlgInterface, \
@@ -34,8 +36,10 @@
     NNHyperoptAlgInterface
 from pytabkit.models.alg_interfaces.xgboost_interfaces import XGBSubSplitInterface, XGBHyperoptAlgInterface, \
     XGBSklearnSubSplitInterface, RandomParamsXGBAlgInterface
+from pytabkit.models.alg_interfaces.xrfm_interfaces import xRFMSubSplitInterface, RandomParamsxRFMAlgInterface
 from pytabkit.models.data.data import TaskType, DictDataset
 from pytabkit.models.nn_models.models import PreprocessingFactory
+from pytabkit.models.torch_utils import TorchTimer
 from pytabkit.models.training.logging import Logger
 from pytabkit.models.training.metrics import Metrics
 
@@ -115,15 +119,13 @@ def run(self, task_package: TaskPackage, logger: Logger, assigned_resources: Nod
 
         interface_resources = assigned_resources.get_interface_resources()
 
-
         old_torch_n_threads = torch.get_num_threads()
         old_torch_n_interop_threads = torch.get_num_interop_threads()
         torch.set_num_threads(interface_resources.n_threads)
         # don't set this because it can throw
         # Error: cannot set number of interop threads after parallel work has started or set_num_interop_threads called
         # torch.set_num_interop_threads(interface_resources.n_threads)
 
-
         ds = task.ds
         name = 'alg ' + task_package.alg_name + ' on task ' + str(task_desc)
 
@@ -185,22 +187,33 @@ def run(self, task_package: TaskPackage, logger: Logger, assigned_resources: Nod
 
         rms = {name: [ResultManager() for _ in task_package.split_infos] for name in pred_param_names}
 
-        cv_alg_interface.fit(ds, cv_idxs_list, interface_resources, logger, cv_tmp_folders, name)
+        with TorchTimer() as cv_fit_timer:
+            cv_alg_interface.fit(ds, cv_idxs_list, interface_resources, logger, cv_tmp_folders, name)
 
         for pred_param_name in pred_param_names:
             cv_alg_interface.set_current_predict_params(pred_param_name)
 
-            cv_results_list = cv_alg_interface.eval(ds, cv_idxs_list, metrics, return_preds)
+            with TorchTimer() as cv_eval_timer:
+                cv_results_list = cv_alg_interface.eval(ds, cv_idxs_list, metrics, return_preds)
 
             for rm, cv_results in zip(rms[pred_param_name], cv_results_list):
-                rm.add_results(is_cv=True, results_dict=cv_results.get_dict())
+                rm.add_results(is_cv=True, results_dict=cv_results.get_dict() |
+                                                        dict(fit_time_s=cv_fit_timer.elapsed,
+                                                             eval_time_s=cv_eval_timer.elapsed))
 
             if n_refit > 0:
                 refit_alg_interface = cv_alg_interface.get_refit_interface(n_refit)
-                refit_results_list = refit_alg_interface.fit_and_eval(ds, refit_idxs_list, interface_resources, logger,
-                                                                      refit_tmp_folders, name, metrics, return_preds)
+
+                with TorchTimer() as refit_fit_timer:
+                    refit_alg_interface.fit(ds, refit_idxs_list, interface_resources, logger, refit_tmp_folders, name)
+
+                with TorchTimer() as refit_eval_timer:
+                    refit_results_list = refit_alg_interface.eval(ds, refit_idxs_list, metrics, return_preds)
                 for rm, refit_results in zip(rms[pred_param_name], refit_results_list):
-                    rm.add_results(is_cv=False, results_dict=refit_results.get_dict())
+                    rm.add_results(is_cv=False,
+                                   results_dict=refit_results.get_dict() |
+                                                dict(fit_time_s=refit_fit_timer.elapsed,
+                                                     eval_time_s=refit_eval_timer.elapsed))
 
         torch.set_num_threads(old_torch_n_threads)
         # torch.set_num_interop_threads(old_torch_n_interop_threads)
@@ -578,3 +591,19 @@ class RandomParamsLinearModelInterfaceWrapper(AlgInterfaceWrapper):
     def __init__(self, model_idx: int, **config):
         # model_idx should be the random search iteration (i.e. start from zero)
         super().__init__(RandomParamsLinearModelAlgInterface, model_idx=model_idx, **config)
+
+
+class TabPFNV2InterfaceWrapper(SubSplitInterfaceWrapper):
+    def create_sub_split_interface(self, task_type: TaskType) -> AlgInterface:
+        return TabPFNV2SubSplitInterface(**self.config)
+
+
+class xRFMInterfaceWrapper(SubSplitInterfaceWrapper):
+    def create_sub_split_interface(self, task_type: TaskType) -> AlgInterface:
+        return xRFMSubSplitInterface(**self.config)
+
+
+class RandomParamsxRFMInterfaceWrapper(MultiSplitAlgInterfaceWrapper):
+    def create_single_alg_interface(self, n_cv: int, task_type: TaskType) \
+            -> AlgInterface:
+        return RandomParamsxRFMAlgInterface(**self.config)

pytabkit/bench/eval/tables.py

@@ -11,14 +11,17 @@
 from pytabkit.models.data.data import TaskType
 from pytabkit.models.data.nested_dict import NestedDict
 
-
-def _get_table_str(table_head: List[List[str]], table_body: List[List[str]]):
-    head_row_strs = [' & '.join(row) + r' \\' for row in table_head]
-    body_row_strs = [' & '.join(row) + r' \\' for row in table_body]
-    n_cols = max(len(row) for row in table_head + table_body)
+def _get_table_str(*parts: List[List[str]]):
+    part_rows = [[' & '.join(row) + r' \\' for row in part] for part in parts]
+    n_cols = max(len(row) for part in parts for row in part)
     begin_table_str = r'\begin{tabular}{' + ('c' * n_cols) + r'}' + '\n' + r'\toprule'
     end_table_str = r'\bottomrule' + '\n' + r'\end{tabular}'
-    all_row_strs = [begin_table_str] + head_row_strs + [r'\midrule'] + body_row_strs + [end_table_str]
+    all_row_strs = [begin_table_str]
+    for part in part_rows[:-1]:
+        all_row_strs.extend(part)
+        all_row_strs.append(r'\midrule')
+    all_row_strs.extend(part_rows[-1])
+    all_row_strs.append(end_table_str)
     complete_str = '\n'.join(all_row_strs)
     return complete_str
 
@@ -208,7 +211,7 @@ def generate_ablations_table(paths: Paths, tables: ResultsTables):
         (r'Activation=SELU', 'act-selu'),
         ('', ''),
         ('No dropout', 'pdrop-0.0'),
-        ('Dropout prob.\ $0.15$ (constant)', 'pdrop-0.15'),
+        (r'Dropout prob.\ $0.15$ (constant)', 'pdrop-0.15'),
         ('', ''),
         ('No weight decay', 'wd-0.0'),
         # ('Weight decay = 0.02 ($\operatorname{flat\_cos}$)', 'wd-0.02-flatcos'),

pytabkit/models/alg_interfaces/alg_interfaces.py

@@ -1,4 +1,5 @@
 import functools
+import warnings
 from pathlib import Path
 from typing import List, Tuple, Any, Optional, Dict
 
@@ -246,6 +247,9 @@ def get_current_predict_params_dict(self):
     def set_current_predict_params(self, name: str) -> None:
         self.curr_pred_params_name = name
 
+    def to(self, device: str) -> None:
+        warnings.warn(f'.to() method does nothing for {self.__class__} (not implemented)')
+
 
 class MultiSplitWrapperAlgInterface(AlgInterface):
     # todo: do we need the option to run this with a "split batch size" > 1 for the NNInterface?

pytabkit/models/alg_interfaces/calibration.py

@@ -41,24 +41,52 @@ def fit(self, ds: DictDataset, idxs_list: List[SplitIdxs], interface_resources:
         self.alg_interface.fit(ds, idxs_list, interface_resources, logger, tmp_folders, name)
         y_preds = self.alg_interface.predict(ds)
 
-        for tt_split_idx, split_idxs in enumerate(idxs_list):
-            for tv_split_idx in range(split_idxs.n_trainval_splits):
-                val_idxs = split_idxs.val_idxs[tv_split_idx]
-                y = ds.tensors['y'][val_idxs]
-                y_pred = y_preds[len(self.calibrators), val_idxs]
-                y_pred_probs = torch.softmax(y_pred, dim=-1)
+        self.n_tv_splits_list_ = [idxs.n_trainval_splits for idxs in idxs_list]
+
+        if self.config.get('calibrate_per_fold', True):
+            for tt_split_idx, split_idxs in enumerate(idxs_list):
+                for tv_split_idx in range(split_idxs.n_trainval_splits):
+                    val_idxs = split_idxs.val_idxs[tv_split_idx]
+                    y = ds.tensors['y'][val_idxs]
+                    y_pred = y_preds[len(self.calibrators), val_idxs]
+                    y_pred_probs = torch.softmax(y_pred, dim=-1)
+
+                    import probmetrics.calibrators
+                    import probmetrics.distributions
+                    calib = probmetrics.calibrators.get_calibrator(**self.config)
+                    if self.config.get('calibrate_with_logits', True):
+                        calib.fit_torch(y_pred=probmetrics.distributions.CategoricalLogits(y_pred.detach().cpu()),
+                                        y_true_labels=y[:, 0])
+                    else:
+                        calib.fit(self._transform_probs(y_pred_probs.detach().cpu().numpy()), y.cpu().numpy()[:, 0])
+
+                    self.calibrators.append(calib)
+                    self.n_calibs.append(val_idxs.shape[-1])
+        else:
+            y_pred_idx = 0
+            for tt_split_idx, split_idxs in enumerate(idxs_list):
+                y_pred_list = []
+                y_list = []
+                for tv_split_idx in range(split_idxs.n_trainval_splits):
+                    val_idxs = split_idxs.val_idxs[tv_split_idx]
+                    y_pred_list.append(y_preds[y_pred_idx, val_idxs])
+                    y_list.append(ds.tensors['y'][val_idxs])
+                    y_pred_idx += 1
+
+                y_pred = torch.cat(y_pred_list, dim=0)
+                y = torch.cat(y_list, dim=0)
 
                 import probmetrics.calibrators
                 import probmetrics.distributions
                 calib = probmetrics.calibrators.get_calibrator(**self.config)
                 if self.config.get('calibrate_with_logits', True):
                     calib.fit_torch(y_pred=probmetrics.distributions.CategoricalLogits(y_pred.detach().cpu()),
-                                    y_true_labels=y[:, 0])
+                                    y_true_labels=y[:, 0].detach().cpu())
                 else:
-                    calib.fit(self._transform_probs(y_pred_probs.detach().cpu().numpy()), y.cpu().numpy()[:, 0])
+                    calib.fit(self._transform_probs(torch.softmax(y_pred, dim=-1).detach().cpu().numpy()), y.cpu().numpy()[:, 0])
 
-                self.calibrators.append(calib)
-                self.n_calibs.append(val_idxs.shape[-1])
+                self.calibrators.extend([calib] * split_idxs.n_trainval_splits)
+                self.n_calibs.extend([y_pred.shape[0]] * split_idxs.n_trainval_splits)
 
         self.fit_params = [dict(sub_fit_params=fp) for fp in self.alg_interface.fit_params]
 
@@ -68,6 +96,15 @@ def predict(self, ds: DictDataset) -> torch.Tensor:
         y_preds = self.alg_interface.predict(ds)
         y_preds_probs = torch.softmax(y_preds, dim=-1)
         y_preds_calib = []
+
+        if self.config.get('ensemble_before_calib', False):
+            start_idx = 0
+            for n_tv_splits in self.n_tv_splits_list_:
+                avg_probs = y_preds_probs[start_idx:start_idx+n_tv_splits].mean(dim=0, keepdim=True)
+                y_preds_probs[start_idx:start_idx + n_tv_splits] = avg_probs
+                start_idx += n_tv_splits
+            y_preds = torch.log(y_preds_probs + 1e-30)
+
         for i in range(y_preds.shape[0]):
             if self.config.get('calibrate_with_logits', True):
                 from probmetrics.distributions import CategoricalLogits
@@ -90,3 +127,7 @@ def predict(self, ds: DictDataset) -> torch.Tensor:
     def get_required_resources(self, ds: DictDataset, n_cv: int, n_refit: int, n_splits: int,
                                split_seeds: List[int], n_train: int) -> RequiredResources:
         return self.alg_interface.get_required_resources(ds, n_cv, n_refit, n_splits, split_seeds, n_train=n_train)
+
+    def to(self, device: str) -> None:
+        self.alg_interface.to(device)
+

pytabkit/models/alg_interfaces/ensemble_interfaces.py

@@ -176,6 +176,12 @@ def get_required_resources(self, ds: DictDataset, n_cv: int, n_refit: int, n_spl
             for ssi in self.alg_interfaces]
         return RequiredResources.combine_sequential(single_resources)
 
+    def to(self, device: str) -> None:
+        for alg_idx, alg_ctx in enumerate(self.alg_contexts_):
+            with alg_ctx as alg_interface:
+                alg_interface.to(device)
+
+
 
 class AlgorithmSelectionAlgInterface(SingleSplitAlgInterface):
     """
@@ -277,6 +283,12 @@ def get_required_resources(self, ds: DictDataset, n_cv: int, n_refit: int, n_spl
             for ssi in self.alg_interfaces]
         return RequiredResources.combine_sequential(single_resources)
 
+    def to(self, device: str) -> None:
+        for alg_idx, alg_ctx in enumerate(self.alg_contexts_):
+            with alg_ctx as alg_interface:
+                alg_interface.to(device)
+
+
 
 class PrecomputedPredictionsAlgInterface(SingleSplitAlgInterface):
     def __init__(self, y_preds_cv: torch.Tensor, y_preds_refit: Optional[torch.Tensor],

pytabkit/models/alg_interfaces/lightgbm_interfaces.py

@@ -123,6 +123,9 @@ def _get_params(self):
                          ('cat_smooth', None),
                          ('cat_l2', None),
                          ('early_stopping_round', ['early_stopping_round', 'early_stopping_rounds'], None),
+                         ('extra_trees', None),
+                         ('max_cat_to_onehot', None),
+                         ('min_data_per_group', None),
                          ]
 
         params = utils.extract_params(self.config, params_config)
@@ -686,10 +689,10 @@ def _sample_params(self, is_classification: bool, seed: int, n_train: int):
                 'min_data_in_leaf': np.floor(np.exp(rng.uniform(np.log(1.0), np.log(65)))),
                 'extra_trees': rng.choice([False, True]),
 
-                'min_data_per_group': np.floor(np.exp(rng.uniform(np.log(2.0), np.log(101)))),
+                'min_data_per_group': round(np.floor(np.exp(rng.uniform(np.log(2.0), np.log(101))))),
                 'cat_l2': np.exp(rng.uniform(np.log(5e-3), np.log(2.0))),
                 'cat_smooth': np.exp(rng.uniform(np.log(1e-3), np.log(100.0))),
-                'max_cat_to_onehot': np.floor(np.exp(rng.uniform(np.log(8.0), np.log(101.0)))),
+                'max_cat_to_onehot': round(np.floor(np.exp(rng.uniform(np.log(8.0), np.log(101.0))))),
 
                 'lambda_l1': np.exp(rng.uniform(np.log(1e-5), np.log(1.0))),
                 'lambda_l2': np.exp(rng.uniform(np.log(1e-5), np.log(2.0))),