fixed device bugs in HPO and ensembling

Author: dholzmueller

README.md

@@ -196,6 +196,7 @@ and https://docs.ray.io/en/latest/cluster/vms/user-guides/community/slurm.html
 
 ## Releases (see git tags)
 
+- v1.5.2: fixed more device bugs for HPO and ensembling
 - v1.5.1: fixed a device bug in TabM for GPU
 - v1.5.0:
     - added `n_repeats` parameter to scikit-learn interfaces for repeated cross-validation

pytabkit/__about__.py

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

pytabkit/models/alg_interfaces/ensemble_interfaces.py

@@ -110,7 +110,7 @@ def fit(self, ds: DictDataset, idxs_list: List[SplitIdxs], interface_resources:
                 weights[weight_idx] += 1
 
                 y_pred_oob = wp.predict_for_weights(weights)
-                loss = Metrics.apply(y_pred_oob, y_oob, val_metric_name).item()
+                loss = Metrics.apply(y_pred_oob.cpu(), y_oob.cpu(), val_metric_name).item()
                 # print(f'{weights=}, {loss=}')
                 if loss < best_step_loss:
                     best_step_loss = loss
@@ -124,7 +124,7 @@ def fit(self, ds: DictDataset, idxs_list: List[SplitIdxs], interface_resources:
                     weights[weight_idx] -= 1
 
                     y_pred_oob = wp.predict_for_weights(weights)
-                    loss = Metrics.apply(y_pred_oob, y_oob, val_metric_name).item()
+                    loss = Metrics.apply(y_pred_oob.cpu(), y_oob.cpu(), val_metric_name).item()
                     # print(f'{weights=}, {loss=}')
                     if loss < best_step_loss:
                         best_step_loss = loss
@@ -233,7 +233,7 @@ def fit(self, ds: DictDataset, idxs_list: List[SplitIdxs], interface_resources:
                 # get out-of-bag predictions
                 y_pred_oob = cat_if_necessary([y_preds[j, idxs_list[0].val_idxs[j]]
                                                for j in range(idxs_list[0].val_idxs.shape[0])], dim=0)
-                loss = Metrics.apply(y_pred_oob, y_oob, val_metric_name).item()
+                loss = Metrics.apply(y_pred_oob.cpu(), y_oob.cpu(), val_metric_name).item()
                 if loss < best_alg_loss:
                     best_alg_loss = loss
                     best_alg_idx = alg_idx

tests/test_variants.py

@@ -3,30 +3,32 @@
 import pandas as pd
 import sklearn
 from sklearn.base import ClassifierMixin
+import torch
 
 from pytabkit import TabM_D_Classifier, RealMLP_HPO_Classifier, Ensemble_HPO_Classifier, TabM_HPO_Regressor, \
     TabM_HPO_Classifier, LGBM_HPO_Classifier, CatBoost_HPO_Classifier, XGB_HPO_Classifier, Ensemble_HPO_Regressor, \
     LGBM_HPO_TPE_Regressor, RealMLP_TD_Regressor, RealMLP_HPO_Regressor, TabM_D_Regressor
 
 
 @pytest.mark.parametrize('estimator', [
-    RealMLP_TD_Regressor(n_cv=2, n_refit=2, n_repeats=2, device='cpu'),
-    RealMLP_HPO_Regressor(device='cpu', n_hyperopt_steps=2, train_metric_name='multi_pinball(0.1,0.9)',
+    RealMLP_TD_Regressor(n_cv=2, n_refit=2, n_repeats=2),
+    RealMLP_HPO_Regressor(n_hyperopt_steps=2, train_metric_name='multi_pinball(0.1,0.9)',
                           val_metric_name='multi_pinball(0.1,0.9)'),
-    TabM_D_Classifier(val_metric_name='cross_entropy', num_emb_type='pwl', tabm_k=16, device='cpu', random_state=0),
-    TabM_D_Regressor(val_metric_name='cross_entropy', num_emb_type='pwl', tabm_k=16, device='cpu', random_state=0),
-    TabM_HPO_Regressor(val_metric_name='mae', n_hyperopt_steps=2, hpo_space_name='tabarena', device='cpu',
+    TabM_D_Classifier(val_metric_name='cross_entropy', num_emb_type='pwl', tabm_k=16, random_state=0),
+    TabM_D_Regressor(val_metric_name='cross_entropy', num_emb_type='pwl', tabm_k=16, random_state=0),
+    TabM_HPO_Regressor(val_metric_name='mae', n_hyperopt_steps=2, hpo_space_name='tabarena',
                        random_state=0),
-    TabM_HPO_Classifier(val_metric_name='mae', n_hyperopt_steps=2, hpo_space_name='default', device='cpu',
+    TabM_HPO_Classifier(val_metric_name='mae', n_hyperopt_steps=2, hpo_space_name='default',
                         random_state=0, use_caruana_ensembling=True),
+    # use CPU since GPU might not support some features in the search space (it has problems with rsm for catboost)
     LGBM_HPO_Classifier(use_caruana_ensembling=True, n_hyperopt_steps=2, hpo_space_name='tabarena', device='cpu'),
     XGB_HPO_Classifier(use_caruana_ensembling=True, n_hyperopt_steps=2, hpo_space_name='tabarena', device='cpu'),
     CatBoost_HPO_Classifier(use_caruana_ensembling=True, n_hyperopt_steps=2, hpo_space_name='tabarena', device='cpu'),
     RealMLP_HPO_Classifier(val_metric_name='cross_entropy', n_hyperopt_steps=3, use_caruana_ensembling=True,
-                           hpo_space_name='tabarena', n_caruana_steps=10, random_state=0, device='cpu'),
-    Ensemble_HPO_Classifier(val_metric_name='brier', device='cpu', n_hpo_steps=2, use_full_caruana_ensembling=True,
+                           hpo_space_name='tabarena', n_caruana_steps=10, random_state=0),
+    Ensemble_HPO_Classifier(val_metric_name='brier', n_hpo_steps=2, use_full_caruana_ensembling=True,
                             use_tabarena_spaces=True),
-    Ensemble_HPO_Regressor(val_metric_name='brier', device='cpu', n_hpo_steps=2, use_full_caruana_ensembling=True,
+    Ensemble_HPO_Regressor(val_metric_name='brier', n_hpo_steps=2, use_full_caruana_ensembling=True,
                            use_tabarena_spaces=True),
     LGBM_HPO_TPE_Regressor(n_cv=2, n_refit=2, n_hyperopt_steps=2),
 ])
@@ -37,7 +39,9 @@ def test_sklearn_not_crash(estimator):
     X['b'] = X['b'].astype('category')
 
     est = sklearn.base.clone(estimator)
-    est.device = 'cpu'
+    if not torch.cuda.is_available():
+        # don't use mps even if it's available
+        est.device = 'cpu'
     if isinstance(est, ClassifierMixin):
         y = np.random.randint(3, size=(n_train,))
     else: