Knowing where time goes keeps experiments fast and avoids oversubscribed CPUs or memory blowups.
scikit-learn uses two levels of parallelism, and they can collide.
- joblib processes or threads, controlled by
n_jobson estimators (forests, k-NN) and oncross_validate/ searches.n_jobs=-1uses all cores. Use it at one level only: a search withn_jobs=-1over a forest withn_jobs=-1oversubscribes. - Native threads in BLAS (NumPy linear algebra) and in OpenMP code such as
HistGradientBoosting. Limit them withthreadpoolctl.threadpool_limits(1)orOMP_NUM_THREADS, especially when joblib already runs many processes. joblib's process backend limits threads in workers automatically.
Other levers:
Pipeline(..., memory="cachedir")caches fitted transformers, so a search that varies only the final model does not refit expensive preprocessing.- Keep data sparse (
OneHotEncoder,TfidfVectorizeroutput) with estimators that support it; densifying text features can use gigabytes. - Use
float32where the estimator preserves it, halving memory. - Pick scalable algorithms:
HistGradientBoosting*instead ofGradientBoosting*,SGDClassifier/SGDRegressororMiniBatchKMeansfor millions of rows. - Out-of-core: estimators with
partial_fit(SGD,MiniBatchKMeans,IncrementalPCA, naive Bayes) learn from chunks streamed from disk; pass allclasseson the first call.
Profile before optimising: often a single slow transformer or an unnecessary dense conversion dominates.
Example
import numpy as np
from sklearn.linear_model import SGDClassifier
from threadpoolctl import threadpool_limits
rng = np.random.default_rng(0)
true_w = rng.normal(size=20)
def chunks(n_chunks=50, size=2000): # e.g. reading a big file in pieces
for _ in range(n_chunks):
X = rng.normal(size=(size, 20)).astype(np.float32)
yield X, (X @ true_w > 0).astype(int)
clf = SGDClassifier(loss="log_loss", random_state=0)
for i, (X, y) in enumerate(chunks()):
clf.partial_fit(X, y, classes=[0, 1] if i == 0 else None)
X_test, y_test = next(chunks(1))
print(round(clf.score(X_test, y_test), 3)) # 0.992
with threadpool_limits(limits=1): # cap BLAS/OpenMP threads
print(clf.predict(X_test[:5])) # [0 0 1 1 0]Remember
- Parallelise at one level: either the search or the estimator
- Cap native threads with threadpoolctl when running many processes
- Use
partial_fitestimators to learn from data that does not fit in memory
Watch out for
- Setting n_jobs=-1 on both GridSearchCV and the RandomForest inside it.