diff --git a/econml/dml/causal_forest.py b/econml/dml/causal_forest.py index d6ff14eeb..d5ac2db2c 100644 --- a/econml/dml/causal_forest.py +++ b/econml/dml/causal_forest.py @@ -683,10 +683,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..f0aa739b4 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: @@ -678,8 +678,10 @@ def _get_inference_options(self): return options def _gen_ortho_learner_model_nuisance(self): - model_propensity = _make_first_stage_selector(self.model_propensity, True, self.random_state) - model_regression = _make_first_stage_selector(self.model_regression, self.discrete_outcome, self.random_state) + model_propensity = _make_first_stage_selector(self.model_propensity, True, self.random_state, + n_jobs=getattr(self, 'n_jobs', None)) + model_regression = _make_first_stage_selector(self.model_regression, self.discrete_outcome, self.random_state, + n_jobs=getattr(self, 'n_jobs', None)) return _ModelNuisance(model_propensity, model_regression, self.min_propensity, self.discrete_outcome) diff --git a/econml/iv/dml/_dml.py b/econml/iv/dml/_dml.py index 12e690e34..a37a33bed 100644 --- a/econml/iv/dml/_dml.py +++ b/econml/iv/dml/_dml.py @@ -409,26 +409,31 @@ def _gen_ortho_learner_model_final(self): return _OrthoIVModelFinal(self._gen_model_final(), self._gen_featurizer(), self.fit_cate_intercept) def _gen_ortho_learner_model_nuisance(self): + n_jobs = getattr(self, 'n_jobs', None) model_y = _make_first_stage_selector(self.model_y_xw, is_discrete=self.discrete_outcome, - random_state=self.random_state) + random_state=self.random_state, + n_jobs=n_jobs) model_t = _make_first_stage_selector(self.model_t_xw, is_discrete=self.discrete_treatment, - random_state=self.random_state) + random_state=self.random_state, + n_jobs=n_jobs) if self.projection: # train E[T|X,W,Z] model_z = _make_first_stage_selector(self.model_t_xwz, is_discrete=self.discrete_treatment, - random_state=self.random_state) + random_state=self.random_state, + n_jobs=n_jobs) else: # train E[Z|X,W] # note: discrete_instrument rather than discrete_treatment in call to _make_first_stage_selector model_z = _make_first_stage_selector(self.model_z_xw, is_discrete=self.discrete_instrument, - random_state=self.random_state) + random_state=self.random_state, + n_jobs=n_jobs) return _OrthoIVNuisanceSelector(model_y, model_t, model_z, self.projection) @@ -1189,13 +1194,16 @@ def _gen_featurizer(self): return clone(self.featurizer, safe=False) def _gen_model_y_xw(self): - return _make_first_stage_selector(self.model_y_xw, self.discrete_outcome, self.random_state) + return _make_first_stage_selector(self.model_y_xw, self.discrete_outcome, self.random_state, + n_jobs=getattr(self, 'n_jobs', None)) def _gen_model_t_xw(self): - return _make_first_stage_selector(self.model_t_xw, self.discrete_treatment, self.random_state) + return _make_first_stage_selector(self.model_t_xw, self.discrete_treatment, self.random_state, + n_jobs=getattr(self, 'n_jobs', None)) def _gen_model_t_xwz(self): - return _make_first_stage_selector(self.model_t_xwz, self.discrete_treatment, self.random_state) + return _make_first_stage_selector(self.model_t_xwz, 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) @@ -1576,13 +1584,16 @@ def _gen_featurizer(self): return clone(self.featurizer, safe=False) def _gen_model_y_xw(self): - return _make_first_stage_selector(self.model_y_xw, self.discrete_outcome, self.random_state) + return _make_first_stage_selector(self.model_y_xw, self.discrete_outcome, self.random_state, + n_jobs=getattr(self, 'n_jobs', None)) def _gen_model_t_xw(self): - return _make_first_stage_selector(self.model_t_xw, self.discrete_treatment, self.random_state) + return _make_first_stage_selector(self.model_t_xw, self.discrete_treatment, self.random_state, + n_jobs=getattr(self, 'n_jobs', None)) def _gen_model_t_xwz(self): - return _make_first_stage_selector(self.model_t_xwz, self.discrete_treatment, self.random_state) + return _make_first_stage_selector(self.model_t_xwz, 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) diff --git a/econml/iv/dr/_dr.py b/econml/iv/dr/_dr.py index b84a79cc7..79961deda 100644 --- a/econml/iv/dr/_dr.py +++ b/econml/iv/dr/_dr.py @@ -663,30 +663,37 @@ def _gen_prel_model_effect(self): return clone(self.prel_model_effect, safe=False) def _gen_ortho_learner_model_nuisance(self): - model_y_xw = _make_first_stage_selector(self.model_y_xw, self.discrete_outcome, self.random_state) - model_t_xw = _make_first_stage_selector(self.model_t_xw, self.discrete_treatment, self.random_state) + n_jobs = getattr(self, 'n_jobs', None) + model_y_xw = _make_first_stage_selector(self.model_y_xw, self.discrete_outcome, self.random_state, + n_jobs=n_jobs) + model_t_xw = _make_first_stage_selector(self.model_t_xw, self.discrete_treatment, self.random_state, + n_jobs=n_jobs) if self.projection: # this is a regression model since the instrument E[T|X,W,Z] is always continuous model_tz_xw = _make_first_stage_selector(self.model_tz_xw, is_discrete=False, - random_state=self.random_state) + random_state=self.random_state, + n_jobs=n_jobs) # we're using E[T|X,W,Z] as the instrument model_z = _make_first_stage_selector(self.model_t_xwz, is_discrete=self.discrete_treatment, - random_state=self.random_state) + random_state=self.random_state, + n_jobs=n_jobs) else: model_tz_xw = _make_first_stage_selector(self.model_tz_xw, is_discrete=(self.discrete_treatment and self.discrete_instrument and not self.fit_cov_directly), - random_state=self.random_state) + random_state=self.random_state, + n_jobs=n_jobs) model_z = _make_first_stage_selector(self.model_z_xw, is_discrete=self.discrete_instrument, - random_state=self.random_state) + random_state=self.random_state, + n_jobs=n_jobs) return [_BaseDRIVNuisanceSelector(prel_model_effect=self._gen_prel_model_effect(), model_y_xw=model_y_xw, @@ -2517,10 +2524,13 @@ def _gen_prel_model_effect(self): return clone(self.prel_model_effect, safe=False) def _gen_ortho_learner_model_nuisance(self): + n_jobs = getattr(self, 'n_jobs', None) model_y_xw = _make_first_stage_selector(self.model_y_xw, is_discrete=self.discrete_outcome, - random_state=self.random_state) - model_t_xwz = _make_first_stage_selector(self.model_t_xwz, is_discrete=True, random_state=self.random_state) + random_state=self.random_state, + n_jobs=n_jobs) + model_t_xwz = _make_first_stage_selector(self.model_t_xwz, is_discrete=True, + random_state=self.random_state, n_jobs=n_jobs) if self.z_propensity == "auto": dummy_z = DummyClassifier(strategy="prior") @@ -2529,7 +2539,8 @@ def _gen_ortho_learner_model_nuisance(self): else: raise ValueError("Only 'auto' or float is allowed!") - dummy_z = _make_first_stage_selector(dummy_z, is_discrete=True, random_state=self.random_state) + dummy_z = _make_first_stage_selector(dummy_z, is_discrete=True, random_state=self.random_state, + n_jobs=n_jobs) return _IntentToTreatDRIVNuisanceSelector(model_y_xw, model_t_xwz, dummy_z, self._gen_prel_model_effect()) diff --git a/econml/panel/dml/_dml.py b/econml/panel/dml/_dml.py index 93a02686b..59e8d4fd8 100644 --- a/econml/panel/dml/_dml.py +++ b/econml/panel/dml/_dml.py @@ -576,12 +576,14 @@ 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 StatsModelsLinearRegression(fit_intercept=False) diff --git a/econml/sklearn_extensions/model_selection.py b/econml/sklearn_extensions/model_selection.py index 8f86a3662..d3f2cfb84 100644 --- a/econml/sklearn_extensions/model_selection.py +++ b/econml/sklearn_extensions/model_selection.py @@ -616,37 +616,38 @@ 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, n_jobs=None, + wrapper=GridSearchCV, needs_scoring=False): named_models = { 'linear': (LogisticRegressionCV(random_state=random_state, cv=cv) if is_discrete - else WeightedLassoCVWrapper(random_state=random_state, cv=cv)), + 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))) + 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(n_jobs=n_jobs, random_state=random_state) if is_discrete + else RandomForestRegressor(n_jobs=n_jobs, random_state=random_state), + 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 return input 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, + random_state=random_state, cv=cv, n_jobs=n_jobs, wrapper=wrapper, needs_scoring=True) # 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, + random_state=random_state, cv=cv, n_jobs=n_jobs, wrapper=wrapper, needs_scoring=needs_scoring) else: raise ValueError(f"Unknown model type: {input}, must be one of {named_models.keys()}")