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..268ea04 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,43 @@ 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: + # Seed only the base learners that have no random_state of their own, + # drawing from self.random_state in this single thread so serial and + # parallel fits agree. Learners that already fix their own randomness are + # left untouched, so their RNG stream is not perturbed (this mirrors how + # scikit learn's forests seed their sub estimators). + seeds = [ + ( + int(self.random_state.randint(np.iinfo(np.int32).max)) + if base.get_params().get("random_state", "set") is None + else None + ) + for base in base_learners + ] + + def _fit_one(base, g, seed): + learner = clone(base) + if seed is not None: + learner.set_params(random_state=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 +638,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))