Performance, parallelism and large data

scikit-learn · Expert · 12 min

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_jobs on estimators (forests, k-NN) and on cross_validate / searches. n_jobs=-1 uses all cores. Use it at one level only: a search with n_jobs=-1 over a forest with n_jobs=-1 oversubscribes.
  • Native threads in BLAS (NumPy linear algebra) and in OpenMP code such as HistGradientBoosting. Limit them with threadpoolctl.threadpool_limits(1) or OMP_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, TfidfVectorizer output) with estimators that support it; densifying text features can use gigabytes.
  • Use float32 where the estimator preserves it, halving memory.
  • Pick scalable algorithms: HistGradientBoosting* instead of GradientBoosting*, SGDClassifier/SGDRegressor or MiniBatchKMeans for millions of rows.
  • Out-of-core: estimators with partial_fit (SGD, MiniBatchKMeans, IncrementalPCA, naive Bayes) learn from chunks streamed from disk; pass all classes on 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_fit estimators 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.

Check yourself

  1. A grid search varies only the classifier's C, but each fit spends most time in a slow TF-IDF step. What helps most?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: scikit-learn.org/stable/computing/parallelism.html