diff --git a/econml/dml/causal_forest.py b/econml/dml/causal_forest.py index e83ca4d4e..2cc3cb79c 100644 --- a/econml/dml/causal_forest.py +++ b/econml/dml/causal_forest.py @@ -680,10 +680,12 @@ def _gen_featurizer(self): return clone(self.featurizer, safe=False) def _gen_model_y(self): - return _make_first_stage_selector(self.model_y, self.discrete_outcome, self.random_state) + return _make_first_stage_selector(self.model_y, self.discrete_outcome, self.random_state, + n_jobs=self.n_jobs) def _gen_model_t(self): - return _make_first_stage_selector(self.model_t, self.discrete_treatment, self.random_state) + return _make_first_stage_selector(self.model_t, self.discrete_treatment, self.random_state, + n_jobs=self.n_jobs) def _gen_model_final(self): return MultiOutputGRF(CausalForest(n_estimators=self.n_estimators, diff --git a/econml/dml/dml.py b/econml/dml/dml.py index 23e47f8ec..cd2f5e079 100644 --- a/econml/dml/dml.py +++ b/econml/dml/dml.py @@ -120,12 +120,13 @@ def best_score(self): return self._model.best_score -def _make_first_stage_selector(model, is_discrete, random_state): +def _make_first_stage_selector(model, is_discrete, random_state, n_jobs=None): if model == 'auto': model = ['forest', 'linear'] return _FirstStageSelector(get_selector(model, is_discrete=is_discrete, - random_state=random_state), + random_state=random_state, + n_jobs=n_jobs), discrete_target=is_discrete) @@ -561,10 +562,12 @@ def _gen_featurizer(self): return clone(self.featurizer, safe=False) def _gen_model_y(self): - return _make_first_stage_selector(self.model_y, self.discrete_outcome, self.random_state) + return _make_first_stage_selector(self.model_y, self.discrete_outcome, self.random_state, + n_jobs=getattr(self, 'n_jobs', None)) def _gen_model_t(self): - return _make_first_stage_selector(self.model_t, self.discrete_treatment, self.random_state) + return _make_first_stage_selector(self.model_t, self.discrete_treatment, self.random_state, + n_jobs=getattr(self, 'n_jobs', None)) def _gen_model_final(self): return clone(self.model_final, safe=False) @@ -1647,11 +1650,13 @@ def _gen_featurizer(self): def _gen_model_y(self): return _make_first_stage_selector(self.model_y, is_discrete=self.discrete_outcome, - random_state=self.random_state) + random_state=self.random_state, + n_jobs=getattr(self, 'n_jobs', None)) def _gen_model_t(self): return _make_first_stage_selector(self.model_t, is_discrete=self.discrete_treatment, - random_state=self.random_state) + random_state=self.random_state, + n_jobs=getattr(self, 'n_jobs', None)) def _gen_model_final(self): return clone(self.model_final, safe=False) diff --git a/econml/dr/_drlearner.py b/econml/dr/_drlearner.py index 256680525..f5e3a9f2a 100644 --- a/econml/dr/_drlearner.py +++ b/econml/dr/_drlearner.py @@ -189,10 +189,10 @@ def predict(self, Y, T, X=None, W=None, *, sample_weight=None, groups=None): return Y_pred.reshape(Y.shape + (T.shape[1] + 1,)), propensities, raw_propensities -def _make_first_stage_selector(model, is_discrete, random_state): +def _make_first_stage_selector(model, is_discrete, random_state, n_jobs=None): if model == "auto": model = ['linear', 'forest'] - return get_selector(model, is_discrete=is_discrete, random_state=random_state) + return get_selector(model, is_discrete=is_discrete, random_state=random_state, n_jobs=n_jobs) class _ModelFinal: diff --git a/econml/sklearn_extensions/model_selection.py b/econml/sklearn_extensions/model_selection.py index 3df538301..8d9e65d74 100644 --- a/econml/sklearn_extensions/model_selection.py +++ b/econml/sklearn_extensions/model_selection.py @@ -437,12 +437,13 @@ def _to_logisticRegression(model: LogisticRegressionCV): _copy_to(model, lr, ["penalty", "dual", "intercept_scaling", "class_weight", "solver", - "verbose", "n_jobs", + "verbose", "tol", "max_iter", "random_state", "n_iter_"]) - # if sklearn version < 1.8, copy multi_class as well + # if sklearn version < 1.8, copy multi_class and n_jobs as well + # (sklearn 1.8 deprecated n_jobs on LogisticRegression; it has no effect post-fit) from packaging import version if version.parse(sklearn.__version__) < version.parse("1.8"): - _copy_to(model, lr, ["multi_class"]) + _copy_to(model, lr, ["multi_class", "n_jobs"]) _copy_to(model, lr, ["classes_"]) _copy_to(model, lr, ["C", "l1_ratio"], True) # these are arrays in LogisticRegressionCV, need to convert them next @@ -616,23 +617,24 @@ def best_score(self): return self._best_score -def get_selector(input, is_discrete, *, random_state=None, cv=None, wrapper=GridSearchCV, needs_scoring=False): +def get_selector(input, is_discrete, *, random_state=None, cv=None, wrapper=GridSearchCV, needs_scoring=False, + n_jobs=None): named_models = { - 'linear': (LogisticRegressionCV(random_state=random_state, cv=cv) if is_discrete - else WeightedLassoCVWrapper(random_state=random_state, cv=cv)), + 'linear': (LogisticRegressionCV(random_state=random_state, cv=cv, n_jobs=n_jobs) if is_discrete + else WeightedLassoCVWrapper(random_state=random_state, cv=cv, n_jobs=n_jobs)), 'poly': ([make_pipeline(PolynomialFeatures(d), - (LogisticRegressionCV(random_state=random_state, cv=cv) if is_discrete - else WeightedLassoCVWrapper(random_state=random_state, cv=cv))) + (LogisticRegressionCV(random_state=random_state, cv=cv, n_jobs=n_jobs) if is_discrete + else WeightedLassoCVWrapper(random_state=random_state, cv=cv, n_jobs=n_jobs))) for d in range(1, 4)]), - 'forest': (GridSearchCV(RandomForestClassifier(random_state=random_state) if is_discrete - else RandomForestRegressor(random_state=random_state), - param_grid={}, cv=cv)), + 'forest': (GridSearchCV(RandomForestClassifier(random_state=random_state, n_jobs=n_jobs) if is_discrete + else RandomForestRegressor(random_state=random_state, n_jobs=n_jobs), + param_grid={}, cv=cv, n_jobs=n_jobs)), 'gbf': (GridSearchCV(GradientBoostingClassifier(random_state=random_state) if is_discrete else GradientBoostingRegressor(random_state=random_state), - param_grid={}, cv=cv)), + param_grid={}, cv=cv, n_jobs=n_jobs)), 'nnet': (GridSearchCV(MLPClassifier(random_state=random_state) if is_discrete else MLPRegressor(random_state=random_state), - param_grid={}, cv=cv)), + param_grid={}, cv=cv, n_jobs=n_jobs)), 'automl': ["poly", "forest", "gbf", "nnet"], } if isinstance(input, ModelSelector): # we've already got a model selector, don't need to do anything @@ -640,14 +642,14 @@ def get_selector(input, is_discrete, *, random_state=None, cv=None, wrapper=Grid elif isinstance(input, list): # we've got a list; call get_selector on each element, then wrap in a ListSelector models = [get_selector(model, is_discrete, random_state=random_state, cv=cv, wrapper=wrapper, - needs_scoring=True) # we need to score to compare outputs to each other + needs_scoring=True, n_jobs=n_jobs) # we need to score to compare outputs to each other for model in input] return ListSelector(models) elif isinstance(input, str): # we've got a string; look it up if input in named_models: return get_selector(named_models[input], is_discrete, random_state=random_state, cv=cv, wrapper=wrapper, - needs_scoring=needs_scoring) + needs_scoring=needs_scoring, n_jobs=n_jobs) else: raise ValueError(f"Unknown model type: {input}, must be one of {named_models.keys()}") elif SklearnCVSelector.can_wrap(input): diff --git a/econml/tests/test_dml.py b/econml/tests/test_dml.py index 2ffb995e9..4b2ba4953 100644 --- a/econml/tests/test_dml.py +++ b/econml/tests/test_dml.py @@ -819,6 +819,55 @@ def true_fn(x): sn6 = est.score_nuisances(Y=y, T=T, X=X, W=W, t_scoring='log_loss') np.testing.assert_allclose(sn6['T_log_loss'], [17.4,17.4], rtol=0, atol=0.1) + def test_n_jobs_propagates_to_first_stage_auto_selector(self): + # Regression test for #1009: SparseLinearDML / CausalForestDML accept + # n_jobs but previously only threaded it into the second-stage Lasso / + # final forest. The 'auto' first-stage selector built RandomForest + + # GridSearchCV + LogisticRegressionCV + WeightedLassoCVWrapper without + # n_jobs, so first-stage fits ran single-core. + def collect_n_jobs(obj, seen=None): + if seen is None: + seen = set() + oid = id(obj) + if oid in seen: + return [] + seen.add(oid) + out = [] + if hasattr(obj, "n_jobs"): + out.append((type(obj).__name__, obj.n_jobs)) + for attr in ("_model", "models", "searcher", "estimator", "_best_model"): + v = getattr(obj, attr, None) + if v is None: + continue + children = v if isinstance(v, (list, tuple)) else [v] + for c in children: + out.extend(collect_n_jobs(c, seen)) + return out + + sentinel = 3 # any non-default int; -1 also works but is harder to assert against + propagating_types = { + 'GridSearchCV', 'RandomForestRegressor', 'RandomForestClassifier', + 'LogisticRegressionCV', 'WeightedLassoCVWrapper', + } + + for est in ( + SparseLinearDML(model_y='auto', model_t='auto', n_jobs=sentinel, + random_state=0), + SparseLinearDML(model_y='auto', model_t='auto', n_jobs=sentinel, + discrete_treatment=True, random_state=0), + CausalForestDML(model_y='auto', model_t='auto', n_jobs=sentinel, + random_state=0), + ): + for selector in (est._gen_model_y(), est._gen_model_t()): + seen_any = False + for name, n_jobs in collect_n_jobs(selector): + if name in propagating_types: + seen_any = True + assert n_jobs == sentinel, \ + f"{type(est).__name__} -> {name}.n_jobs = {n_jobs}, expected {sentinel}" + assert seen_any, \ + f"selector tree for {type(est).__name__} had no n_jobs-bearing leaves" + def test_aaforest_pandas(self): """Test that we can use CausalForest with pandas inputs.""" df = pd.DataFrame({'a': np.random.normal(size=500),