From 886522f7999fb35c2dc6ed3aabe71943af61acd5 Mon Sep 17 00:00:00 2001 From: loganprosser Date: Wed, 29 Jul 2026 20:52:09 -0400 Subject: [PATCH] Add optional n_jobs to fit base learners in parallel NGBoost fits one base learner per distribution parameter on every boosting round. For distributions with many parameters, such as MultivariateNormal, those independent fits run one after another and take most of the training time. This adds an optional n_jobs argument. The default (None) fits serially; a value above 1, or -1 for all cores, fits the base learners at the same time with joblib's threading backend, giving roughly a 2x speedup on high dimensional MultivariateNormal in local benchmarks. scikit learn releases the GIL while building a tree, so threads parallelize without copying X to worker processes each round, and the backend is set to threading explicitly so an outer parallel context such as a parallel GridSearchCV cannot turn the inner fits into processes. Each base learner is seeded from the model random_state before dispatch (only when the learner has no random_state of its own), on both the serial and parallel paths, so the fitted model is independent of n_jobs and reproducible when random_state is set. This mirrors how scikit learn's forests seed trees. It does shift the default base learner's output relative to 0.5.x, noted in RELEASE_NOTES. For compressed sparse inputs (CSR, CSC, BSR) the shared matrix is canonicalized once with sort_indices before the threaded fits so they cannot race on an in place index sort; other sparse formats are copied by scikit learn. NGBRegressor, NGBClassifier and NGBSurvival forward n_jobs and get_params exposes it. joblib is listed in pyproject (already present via scikit learn). Tests cover get_params and clone, and serial versus parallel equivalence on regression, multiclass classification, sparse float32 CSC input, missing values, and the default base learner. This change was authored with assistance from an AI coding agent. --- RELEASE_NOTES.md | 1 + ngboost/api.py | 13 ++++- ngboost/ngboost.py | 50 ++++++++++++++++++-- pyproject.toml | 1 + tests/test_basic.py | 113 +++++++++++++++++++++++++++++++++++++++++++- 5 files changed, 170 insertions(+), 8 deletions(-) diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index 4ce1721..22f2068 100644 --- a/RELEASE_NOTES.md +++ b/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) diff --git a/ngboost/api.py b/ngboost/api.py index 84d5413..2051e59 100644 --- a/ngboost/api.py +++ b/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): diff --git a/ngboost/ngboost.py b/ngboost/ngboost.py index 2bae7a2..ad7abc4 100644 --- a/ngboost/ngboost.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index fbd62cc..eea34b9 100644 --- a/pyproject.toml +++ b/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'"} diff --git a/tests/test_basic.py b/tests/test_basic.py index 2f65ff1..ae9c8df 100644 --- a/tests/test_basic.py +++ b/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))