ENH: move _partition_estimators to ensemble.base

Author: glouppe

sklearn/ensemble/bagging.py

@@ -13,7 +13,7 @@
 from inspect import getargspec
 
 from ..base import ClassifierMixin, RegressorMixin
-from ..externals.joblib import Parallel, delayed, cpu_count
+from ..externals.joblib import Parallel, delayed
 from ..externals.six import with_metaclass
 from ..externals.six.moves import zip
 from ..metrics import r2_score, accuracy_score
@@ -22,7 +22,7 @@
 from ..utils.fixes import bincount, unique
 from ..utils.random import sample_without_replacement
 
-from .base import BaseEnsemble
+from .base import BaseEnsemble, _partition_estimators
 
 __all__ = ["BaggingClassifier",
            "BaggingRegressor"]
@@ -186,24 +186,6 @@ def _parallel_predict_regression(estimators, estimators_features, X):
                                               estimators_features))
 
 
-def _partition_estimators(ensemble):
-    """Private function used to partition estimators between jobs."""
-    # Compute the number of jobs
-    if ensemble.n_jobs == -1:
-        n_jobs = min(cpu_count(), ensemble.n_estimators)
-
-    else:
-        n_jobs = min(ensemble.n_jobs, ensemble.n_estimators)
-
-    # Partition estimators between jobs
-    n_estimators = (ensemble.n_estimators // n_jobs) * np.ones(n_jobs,
-                                                               dtype=np.int)
-    n_estimators[:ensemble.n_estimators % n_jobs] += 1
-    starts = np.cumsum(n_estimators)
-
-    return n_jobs, n_estimators.tolist(), [0] + starts.tolist()
-
-
 class BaseBagging(with_metaclass(ABCMeta, BaseEnsemble)):
     """Base class for Bagging meta-estimator.
 

sklearn/ensemble/base.py

@@ -5,9 +5,12 @@
 # Authors: Gilles Louppe
 # License: BSD 3 clause
 
+import numpy as np
+
 from ..base import clone
 from ..base import BaseEstimator
 from ..base import MetaEstimatorMixin
+from ..externals.joblib import cpu_count
 
 
 class BaseEnsemble(BaseEstimator, MetaEstimatorMixin):
@@ -78,3 +81,21 @@ def __getitem__(self, index):
     def __iter__(self):
         """Returns iterator over estimators in the ensemble."""
         return iter(self.estimators_)
+
+
+def _partition_estimators(ensemble):
+    """Private function used to partition estimators between jobs."""
+    # Compute the number of jobs
+    if ensemble.n_jobs == -1:
+        n_jobs = min(cpu_count(), ensemble.n_estimators)
+
+    else:
+        n_jobs = min(ensemble.n_jobs, ensemble.n_estimators)
+
+    # Partition estimators between jobs
+    n_estimators = (ensemble.n_estimators // n_jobs) * np.ones(n_jobs,
+                                                               dtype=np.int)
+    n_estimators[:ensemble.n_estimators % n_jobs] += 1
+    starts = np.cumsum(n_estimators)
+
+    return n_jobs, n_estimators.tolist(), [0] + starts.tolist()

sklearn/ensemble/forest.py

@@ -44,7 +44,7 @@ class calls the ``fit`` method of each sub-estimator on random samples
 from abc import ABCMeta, abstractmethod
 
 from ..base import ClassifierMixin, RegressorMixin
-from ..externals.joblib import Parallel, delayed, cpu_count
+from ..externals.joblib import Parallel, delayed
 from ..externals import six
 from ..externals.six.moves import xrange
 from ..feature_selection.from_model import _LearntSelectorMixin
@@ -57,8 +57,7 @@ class calls the ``fit`` method of each sub-estimator on random samples
 from ..utils.validation import DataConversionWarning
 from ..utils.fixes import bincount, unique
 
-
-from .base import BaseEnsemble
+from .base import BaseEnsemble, _partition_estimators
 
 __all__ = ["RandomForestClassifier",
            "RandomForestRegressor",
@@ -151,29 +150,6 @@ def _parallel_predict_regression(trees, X):
     return sum(tree.predict(X) for tree in trees)
 
 
-def _partition_trees(forest):
-    """Private function used to partition trees between jobs."""
-    # Compute the number of jobs
-    if forest.n_jobs == -1:
-        n_jobs = min(cpu_count(), forest.n_estimators)
-
-    else:
-        n_jobs = min(forest.n_jobs, forest.n_estimators)
-
-    # Partition trees between jobs
-    n_trees = [forest.n_estimators // n_jobs] * n_jobs
-
-    for i in range(forest.n_estimators % n_jobs):
-        n_trees[i] += 1
-
-    starts = [0] * (n_jobs + 1)
-
-    for i in range(1, n_jobs + 1):
-        starts[i] = starts[i - 1] + n_trees[i - 1]
-
-    return n_jobs, n_trees, starts
-
-
 class BaseForest(six.with_metaclass(ABCMeta, BaseEnsemble,
                                     _LearntSelectorMixin)):
     """Base class for forests of trees.
@@ -286,7 +262,7 @@ def fit(self, X, y, sample_weight=None):
                              " if bootstrap=True")
 
         # Assign chunk of trees to jobs
-        n_jobs, n_trees, _ = _partition_trees(self)
+        n_jobs, n_trees, _ = _partition_estimators(self)
 
         # Precalculate the random states
         seeds = [random_state.randint(MAX_INT, size=i) for i in n_trees]
@@ -481,7 +457,7 @@ def predict_proba(self, X):
             X = array2d(X, dtype=DTYPE)
 
         # Assign chunk of trees to jobs
-        n_jobs, n_trees, starts = _partition_trees(self)
+        n_jobs, n_trees, starts = _partition_estimators(self)
 
         # Parallel loop
         all_proba = Parallel(n_jobs=n_jobs, verbose=self.verbose)(
@@ -590,7 +566,7 @@ def predict(self, X):
             X = array2d(X, dtype=DTYPE)
 
         # Assign chunk of trees to jobs
-        n_jobs, n_trees, starts = _partition_trees(self)
+        n_jobs, n_trees, starts = _partition_estimators(self)
 
         # Parallel loop
         all_y_hat = Parallel(n_jobs=n_jobs, verbose=self.verbose)(