[ENH] Add configurable predictor parameters (LAMDA-NJU#10)

Author: tczhao

CHANGELOG.rst

@@ -30,3 +30,5 @@ Version 0.1.*
 .. |Enhancement| replace:: :raw-html:`<span class="badge badge-info">Enhancement</span>` :raw-latex:`{\small\sc [Enhancement]}`
 .. |Fix| replace:: :raw-html:`<span class="badge badge-danger">Fix</span>` :raw-latex:`{\small\sc [Fix]}`
 .. |API| replace:: :raw-html:`<span class="badge badge-warning">API Change</span>` :raw-latex:`{\small\sc [API Change]}`
+
+- |Feature| configurable predictor parameter `#9 <https://github.com/LAMDA-NJU/Deep-Forest/issues/10>`__

deepforest/cascade.py

@@ -14,14 +14,23 @@
 from ._binner import Binner
 
 
+def _get_predictor_kwargs(predictor_kwargs, **kwargs) -> dict:
+    """Overwrites default args if predictor_kwargs is supplied."""
+    for key, value in kwargs.items():
+        if key not in predictor_kwargs.keys():
+            predictor_kwargs[key] = value
+    return predictor_kwargs
+
+
 def _build_predictor(
     predictor_name,
     n_estimators,
     n_outputs,
     max_depth=None,
     min_samples_leaf=1,
     n_jobs=None,
-    random_state=None
+    random_state=None,
+    predictor_kwargs={},
 ):
     """Build the predictor concatenated to the deep forest."""
     predictor_name = predictor_name.lower()
@@ -30,11 +39,14 @@ def _build_predictor(
     if predictor_name == "forest":
         from .forest import RandomForestClassifier
         predictor = RandomForestClassifier(
-            n_estimators=n_estimators,
-            max_depth=max_depth,
-            min_samples_leaf=min_samples_leaf,
-            n_jobs=n_jobs,
-            random_state=random_state,
+            **_get_predictor_kwargs(
+                predictor_kwargs,
+                n_estimators=n_estimators,
+                max_depth=max_depth,
+                min_samples_leaf=min_samples_leaf,
+                n_jobs=n_jobs,
+                random_state=random_state,
+            )
         )
     # XGBoost
     elif predictor_name == "xgboost":
@@ -51,11 +63,14 @@ def _build_predictor(
         # because the exact mode of XGBoost is too slow.
         objective = "multi:softmax" if n_outputs > 2 else "binary:logistic"
         predictor = xgb.sklearn.XGBClassifier(
-            objective=objective,
-            n_estimators=n_estimators,
-            tree_method="hist",
-            n_jobs=n_jobs,
-            random_state=random_state,
+            **_get_predictor_kwargs(
+                predictor_kwargs,
+                objective=objective,
+                n_estimators=n_estimators,
+                tree_method="hist",
+                n_jobs=n_jobs,
+                random_state=random_state,
+            )
         )
     # LightGBM
     elif predictor_name == "lightgbm":
@@ -70,10 +85,13 @@ def _build_predictor(
 
         objective = "multiclass" if n_outputs > 2 else "binary"
         predictor = lgb.LGBMClassifier(
-            objective=objective,
-            n_estimators=n_estimators,
-            n_jobs=n_jobs,
-            random_state=random_state,
+            **_get_predictor_kwargs(
+                predictor_kwargs,
+                objective=objective,
+                n_estimators=n_estimators,
+                n_jobs=n_jobs,
+                random_state=random_state,
+            )
         )
     else:
         msg = (
@@ -117,6 +135,11 @@ def _build_predictor(
     predictor : :obj:`{"forest", "xgboost", "lightgbm"}`, default="forest"
         The type of the predictor concatenated to the deep forest. If
         ``use_predictor`` is False, this parameter will have no effect.
+    predictor_kwargs : :obj:`dict`, default={}
+        The configuration of the predictor concatenated to the deep forest.
+        Specifying this will extend/overwrite the original parameters inherit
+        from deep forest.
+        If ``use_predictor`` is False, this parameter will have no effect.
     n_tolerant_rounds : :obj:`int`, default=2
         Specify when to conduct early stopping. The training process
         terminates when the validation performance on the training set does
@@ -182,6 +205,7 @@ def __init__(
         min_samples_leaf=1,
         use_predictor=False,
         predictor="forest",
+        predictor_kwargs={},
         n_tolerant_rounds=2,
         delta=1e-5,
         partial_mode=False,
@@ -197,6 +221,7 @@ def __init__(
         self.n_trees = n_trees
         self.max_depth = max_depth
         self.min_samples_leaf = min_samples_leaf
+        self.predictor_kwargs = predictor_kwargs
         self.n_tolerant_rounds = n_tolerant_rounds
         self.delta = delta
         self.partial_mode = partial_mode
@@ -618,7 +643,8 @@ def fit(self, X, y):
                 self.max_depth,
                 self.min_samples_leaf,
                 self.n_jobs,
-                self.random_state
+                self.random_state,
+                self.predictor_kwargs,
             )
 
             binner_ = Binner(

tests/test_model.py

@@ -7,6 +7,7 @@
 
 import deepforest
 from deepforest import CascadeForestClassifier
+from deepforest.cascade import _get_predictor_kwargs
 
 
 save_dir = "./tmp"
@@ -26,6 +27,7 @@
               "min_samples_leaf": 1,
               "use_predictor": True,
               "predictor": "forest",
+              "predictor_kwargs": {},
               "n_tolerant_rounds": 2,
               "delta": 1e-5,
               "n_jobs": -1,
@@ -41,13 +43,35 @@
           "min_samples_leaf": 1,
           "use_predictor": True,
           "predictor": "forest",
+          "predictor_kwargs": {},
           "n_tolerant_rounds": 2,
           "delta": 1e-5,
           "n_jobs": -1,
           "random_state": 0,
           "verbose": 2}
 
 
+@pytest.mark.parametrize(
+    "test_input,expected",
+    [
+        (
+            {"predictor_kwargs": {}, "n_job": 2},
+            {"n_job": 2},
+        ),
+        (
+            {"predictor_kwargs": {"n_job": 3}, "n_job": 2},
+            {"n_job": 3},
+        ),
+        (
+            {"predictor_kwargs": {"iter": 4}, "n_job": 2},
+            {"iter": 4, "n_job": 2},
+        ),
+    ],
+)
+def test_predictor_kwargs_overwrite(test_input, expected):
+    assert _get_predictor_kwargs(**test_input) == expected
+
+
 def test_model_properties_after_fitting():
     """Check the model properties after fitting a deep forest model."""
     model = CascadeForestClassifier(**toy_kwargs)