Fix n_jobs not propagated to auto first-stage learners - #1045
Open
genrichez wants to merge 2 commits into
Open
Conversation
When using model_y='auto' or model_t='auto', the n_jobs parameter was not forwarded to the first-stage model selector. This meant that RandomForest and other parallelizable models in the auto selector always ran on a single core regardless of the n_jobs setting. Add n_jobs parameter to get_selector() and _make_first_stage_selector(), and forward it from all DML, DR, and IV estimator classes that have n_jobs. Classes without n_jobs (e.g. LinearDML) default to None. Skip LogisticRegressionCV since sklearn 1.8 deprecated its n_jobs param. Fixes microsoft#1009 Signed-off-by: genrichez <2.2434764e+07+genrichez@users.noreply.github.com>
Author
|
Hi @kbattocchi following up on this one since it's been quiet a few weeks. You folded my other PR #1046 into #1047 recently, so figured I'd flag this one too in case it's useful, fixes n_jobs not propagating to the auto first-stage learners. Happy to adjust if you'd like a different approach. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
When using model_y='auto' or model_t='auto' in estimators like SparseLinearDML or CausalForestDML, the n_jobs parameter was not forwarded to the first-stage model selector. This meant that RandomForest and other parallelizable models always ran on a single core, regardless of the n_jobs setting passed to the estimator.
Fixes #1009
Changes
Added n_jobs parameter to get_selector( ) and _make_first_stage_selector(), and forwarded it from all estimator classes that support n_jobs:
econml/sklearn_extensions/model_selection.pyget_selector()now acceptsn_jobsand passes it to RF, GridSearchCV, and WeightedLassoCVWrappereconml/dml/dml.py_make_first_stage_selector()acceptsn_jobs;DMLandNonParamDMLforward iteconml/dml/causal_forest.pyCausalForestDMLforwardsself.n_jobseconml/dr/_drlearner.py_make_first_stage_selector()updatedeconml/iv/dml/_dml.pyn_jobseconml/iv/dr/_dr.pyn_jobseconml/panel/dml/_dml.pyDynamicDMLforwardsn_jobsClasses without n_jobs (e.g. LinearDML) use getattr(self, 'n_jobs', None) so they default to None with no behavior change.
LogisticRegressionCV is intentionally skipped because sklearn 1.8 deprecated its n_jobs parameter.
How to verify
Tests
All existing tests pass (excluding unrelated ray tests that require the ray package)