Add optional n_jobs to fit base learners in parallel

RELEASE_NOTES.md

@@ -2,6 +2,7 @@
 
 ## Version 0.5.11
 
+* Added an optional n_jobs parameter that fits the per parameter base learners in parallel with threads. Base learners are now seeded from the model random_state, so fits are reproducible and independent of n_jobs; output for the default base learner shifts relative to 0.5.x.
 * Fix `NGBClassifier` and `NGBSurvival` API parity by adding `validation_fraction` and `early_stopping_rounds`, with regression coverage for their validation-split paths (issue #402)
 * Replace the deprecated `friedman_mse` criterion with `squared_error` for the default tree learner and distribution tests (issue #408, PR #409)
 * Support a distinct base learner for each distribution parameter, including scikit-learn nested parameters and feature importances for all-tree configurations (issue #338, PR #411)

ngboost/api.py

@@ -57,12 +57,13 @@ class NGBRegressor(NGBoost, BaseEstimator):
                                     loss has to increase before the algorithm stops early.
                                     Set to None to disable early stopping and validation.
                                     None enables running over the full data set.
+        n_jobs            : see the ``NGBoost`` base class.
 
     Output:
         An NGBRegressor object that can be fit.
     """
 
-    # pylint: disable=too-many-positional-arguments
+    # pylint: disable=too-many-positional-arguments,too-many-locals
     def __init__(
         self,
         Dist=Normal,
@@ -79,6 +80,7 @@ def __init__(
         random_state=None,
         validation_fraction=0.1,
         early_stopping_rounds=None,
+        n_jobs=None,
     ):
         assert issubclass(
             Dist, RegressionDistn
@@ -104,6 +106,7 @@ def __init__(
             random_state,
             validation_fraction,
             early_stopping_rounds,
+            n_jobs,
         )
 
         self._estimator_type = "regressor"
@@ -160,11 +163,12 @@ class NGBClassifier(NGBoost, BaseEstimator):
                                     loss has to increase before the algorithm stops early.
                                     Set to None to disable early stopping and validation.
                                     None enables running over the full data set.
+        n_jobs            : see the ``NGBoost`` base class.
     Output:
         An NGBClassifier object that can be fit.
     """
 
-    # pylint: disable=too-many-positional-arguments
+    # pylint: disable=too-many-positional-arguments,too-many-locals
     def __init__(
         self,
         Dist=Bernoulli,
@@ -181,6 +185,7 @@ def __init__(
         random_state=None,
         validation_fraction=0.1,
         early_stopping_rounds=None,
+        n_jobs=None,
     ):
         assert issubclass(
             Dist, ClassificationDistn
@@ -200,6 +205,7 @@ def __init__(
             random_state,
             validation_fraction,
             early_stopping_rounds,
+            n_jobs,
         )
         self._estimator_type = "classifier"
 
@@ -271,6 +277,7 @@ class NGBSurvival(NGBoost, BaseEstimator):
                                     loss has to increase before the algorithm stops early.
                                     Set to None to disable early stopping and validation.
                                     None enables running over the full data set.
+        n_jobs            : see the ``NGBoost`` base class.
     Output:
         An NGBSurvival object that can be fit.
     """
@@ -292,6 +299,7 @@ def __init__(
         random_state=None,
         validation_fraction=0.1,
         early_stopping_rounds=None,
+        n_jobs=None,
     ):
 
         assert issubclass(
@@ -320,6 +328,7 @@ def __init__(
             random_state,
             validation_fraction,
             early_stopping_rounds,
+            n_jobs,
         )
 
     def __getstate__(self):

ngboost/ngboost.py

@@ -5,6 +5,8 @@
 # pylint: disable=unused-variable,invalid-unary-operand-type,attribute-defined-outside-init
 # pylint: disable=redundant-keyword-arg,protected-access,unnecessary-lambda-assignment
 import numpy as np
+from joblib import Parallel, delayed
+from scipy.sparse import issparse
 from sklearn.base import clone
 from sklearn.model_selection import train_test_split
 from sklearn.tree import DecisionTreeRegressor
@@ -53,6 +55,18 @@ class NGBoost:
                                     loss has to increase before the algorithm stops early.
                                     Set to None to disable early stopping and validation.
                                     None enables running over the full data set.
+        n_jobs            : number of joblib threads used to fit the per parameter
+                            base learners in parallel on each boosting round. None
+                            or 1 (the default) fits them serially. A value above 1,
+                            or ``-1`` for all cores, speeds up fitting when the
+                            distribution has many parameters, such as a high
+                            dimensional MultivariateNormal or classification with
+                            many classes; on small problems the thread overhead can
+                            make it slower instead. Threads are used so the shared
+                            data is not copied per worker; when nesting inside an
+                            outer parallel job keep the product of the thread counts
+                            near the core count. The fitted model does not depend on
+                            n_jobs.
 
 
     Output:
@@ -76,6 +90,7 @@ def __init__(
         random_state=None,
         validation_fraction=0.1,
         early_stopping_rounds=None,
+        n_jobs=None,
     ):
         self.Dist = Dist
         self.Score = Score
@@ -98,6 +113,7 @@ def __init__(
         self.best_val_loss_itr = None
         self.validation_fraction = validation_fraction
         self.early_stopping_rounds = early_stopping_rounds
+        self.n_jobs = n_jobs
 
         if hasattr(self.Dist, "multi_output"):
             self.multi_output = self.Dist.multi_output
@@ -189,13 +205,36 @@ def _base_learners(self):
 
     def fit_base(self, X, grads, sample_weight=None):
         base_learners = self._base_learners()
-        if sample_weight is None:
-            models = [clone(base).fit(X, g) for base, g in zip(base_learners, grads.T)]
-        else:
+        # Draw one seed per base learner from self.random_state in this single
+        # thread, so serial and parallel fits agree and results are reproducible
+        # when random_state is set (mirrors how scikit-learn's forests seed trees).
+        seeds = self.random_state.randint(
+            np.iinfo(np.int32).max, size=len(base_learners)
+        )
+
+        def _fit_one(base, g, seed):
+            learner = clone(base)
+            if learner.get_params().get("random_state", "unset") is None:
+                learner.set_params(random_state=int(seed))
+            if sample_weight is None:
+                return learner.fit(X, g)
+            return learner.fit(X, g, sample_weight=sample_weight)
+
+        if self.n_jobs in (None, 1):
             models = [
-                clone(base).fit(X, g, sample_weight=sample_weight)
-                for base, g in zip(base_learners, grads.T)
+                _fit_one(base, g, seed)
+                for base, g, seed in zip(base_learners, grads.T, seeds)
             ]
+        else:
+            if issparse(X) and X.format in ("csr", "csc", "bsr"):
+                # Canonicalize once so the threaded fits do not race on an in
+                # place sort_indices of this shared matrix. Only the compressed
+                # formats have sortable indices; sklearn copies the rest.
+                X.sort_indices()
+            models = Parallel(n_jobs=self.n_jobs, backend="threading")(
+                delayed(_fit_one)(base, g, seed)
+                for base, g, seed in zip(base_learners, grads.T, seeds)
+            )
         fitted = np.array([m.predict(X) for m in models]).T
         self.base_models.append(models)
         return fitted
@@ -592,6 +631,7 @@ def get_params(self, deep=True):
             "random_state": self.random_state,
             "validation_fraction": self.validation_fraction,
             "early_stopping_rounds": self.early_stopping_rounds,
+            "n_jobs": self.n_jobs,
         }
 
         if not deep:

pyproject.toml

@@ -29,6 +29,7 @@ scipy = [
 tqdm = ">=4.3"
 lifelines = ">=0.25"
 sympy = ">=1.12"
+joblib = ">=1.2"
 matplotlib = [
     {version = ">=3.0,<3.10", markers = "python_version < '3.10'"},
     {version = ">=3.10.0", markers = "python_version >= '3.10'"}

tests/test_basic.py

@@ -1,12 +1,13 @@
 import numpy as np
 import pytest
+from scipy import sparse
 from sklearn.base import BaseEstimator, RegressorMixin, clone
 from sklearn.ensemble import HistGradientBoostingRegressor
 from sklearn.linear_model import Ridge
 from sklearn.tree import DecisionTreeRegressor
 
 from ngboost import NGBClassifier, NGBRegressor, NGBSurvival
-from ngboost.distns import Bernoulli, Normal, k_categorical
+from ngboost.distns import Bernoulli, MultivariateNormal, Normal, k_categorical
 
 
 class RecordingRegressor(BaseEstimator, RegressorMixin):
@@ -351,3 +352,113 @@ def test_feature_importances_are_none_for_mixed_base_learners():
     ).fit(X, Y)
 
     assert ngb.feature_importances_ is None
+
+
+def test_n_jobs_is_a_supported_param():
+    ngb = NGBRegressor(n_jobs=4, verbose=False)
+    assert ngb.get_params()["n_jobs"] == 4
+    assert clone(ngb).get_params()["n_jobs"] == 4
+
+
+def test_parallel_tree_fits_match_serial():
+    rng = np.random.default_rng(0)
+    X = rng.normal(size=(200, 5))
+    W = rng.normal(size=(5, 3))
+    Y = X @ W + 0.1 * rng.normal(size=(200, 3))
+
+    kw = {
+        "Dist": MultivariateNormal(3),
+        "Base": DecisionTreeRegressor(max_depth=3, random_state=0),
+        "n_estimators": 10,
+        "verbose": False,
+    }
+    serial = NGBRegressor(n_jobs=1, **kw).fit(X, Y)
+    parallel = NGBRegressor(n_jobs=-1, **kw).fit(X, Y)
+
+    np.testing.assert_allclose(serial.pred_param(X), parallel.pred_param(X))
+
+
+def test_parallel_matches_serial_default_random_state():
+    # With the model random_state set, per learner seeding makes the default
+    # random_state=None base learner reproducible and independent of n_jobs.
+    rng = np.random.default_rng(3)
+    X = rng.normal(size=(200, 5))
+    W = rng.normal(size=(5, 3))
+    Y = X @ W + 0.1 * rng.normal(size=(200, 3))
+
+    kw = {"Dist": MultivariateNormal(3), "n_estimators": 10, "verbose": False}
+    serial = NGBRegressor(random_state=0, n_jobs=1, **kw).fit(X, Y)
+    parallel = NGBRegressor(random_state=0, n_jobs=-1, **kw).fit(X, Y)
+
+    np.testing.assert_allclose(serial.pred_param(X), parallel.pred_param(X))
+
+
+def test_parallel_matches_serial_sparse_float32_csc():
+    # A float32 CSC matrix with unsorted indices is the exact input where a
+    # threaded in place sort_indices could race across the per parameter fits.
+    # fit_base canonicalizes the shared matrix once first, so parallel must
+    # still match serial. Repeat a few times to make a regression obvious.
+    rng = np.random.default_rng(0)
+    Xd = rng.normal(size=(300, 8)).astype(np.float32)
+    Xd[Xd < 0.6] = 0.0
+    Y = rng.integers(0, 6, size=300)
+
+    def make_X():
+        X = sparse.csc_matrix(Xd)
+        X.sort_indices()
+        for j in range(X.shape[1]):
+            s, e = X.indptr[j], X.indptr[j + 1]
+            if e - s > 1:
+                perm = rng.permutation(e - s)
+                X.indices[s:e] = X.indices[s:e][perm]
+                X.data[s:e] = X.data[s:e][perm]
+        X.has_sorted_indices = False
+        return X
+
+    kw = {
+        "Dist": k_categorical(6),
+        "Base": DecisionTreeRegressor(max_depth=3, random_state=0),
+        "n_estimators": 8,
+        "verbose": False,
+    }
+    ref = NGBClassifier(n_jobs=1, **kw).fit(make_X(), Y).pred_param(Xd)
+    for _ in range(5):
+        got = NGBClassifier(n_jobs=-1, **kw).fit(make_X(), Y).pred_param(Xd)
+        np.testing.assert_allclose(got, ref)
+
+
+def test_parallel_matches_serial_multiclass():
+    rng = np.random.default_rng(1)
+    X = rng.normal(size=(150, 4))
+    Y = np.argmax(np.column_stack([X[:, 0], X[:, 1], -X[:, 0] - X[:, 1]]), axis=1)
+
+    kw = {
+        "Dist": k_categorical(3),
+        "Base": DecisionTreeRegressor(max_depth=2, random_state=0),
+        "n_estimators": 10,
+        "verbose": False,
+    }
+    serial = NGBClassifier(n_jobs=1, **kw).fit(X, Y)
+    parallel = NGBClassifier(n_jobs=-1, **kw).fit(X, Y)
+
+    np.testing.assert_allclose(serial.predict_proba(X), parallel.predict_proba(X))
+
+
+def test_parallel_matches_serial_with_missing_values():
+    # With a fixed random_state on the base learner, parallel fitting stays
+    # deterministic and matches serial even on the missing value path, which
+    # otherwise races on numpy's global random state under threads.
+    rng = np.random.default_rng(2)
+    X = rng.normal(size=(200, 5))
+    X[rng.random(X.shape) < 0.1] = np.nan
+    Y = rng.normal(size=200)
+
+    kw = {
+        "Base": DecisionTreeRegressor(max_depth=3, random_state=0),
+        "n_estimators": 10,
+        "verbose": False,
+    }
+    serial = NGBRegressor(n_jobs=1, **kw).fit(X, Y)
+    parallel = NGBRegressor(n_jobs=-1, **kw).fit(X, Y)
+
+    np.testing.assert_allclose(serial.pred_param(X), parallel.pred_param(X))