Use this with https://github.com/scikit-learn/scikit-learn/pull/34667 (tested with `16de9236116f0dd3dcc769eaea781d389a6ae419`)

bench_store_fitted_as_numpy.py

"""Benchmark storing fitted attributes as NumPy.

Times a PCA + LogisticRegression pipeline with and without the
`store_fitted_as_numpy` configuration option added in
https://github.com/scikit-learn/scikit-learn/pull/34667, for
https://github.com/scikit-learn/scikit-learn/issues/34604.

Needs that branch checked out. Run it with no arguments:

    python bench_store_fitted_as_numpy.py

Edit DEVICES below to pick which namespace and device combinations to run;
whatever is not installed or not available is skipped.
"""

# isort: skip_file

import os
os.environ.setdefault("SCIPY_ARRAY_API", "1")
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")

import importlib.util
import time
import warnings

import numpy as np

from sklearn import config_context
from sklearn.datasets import make_classification
from sklearn.decomposition import PCA
from sklearn.exceptions import ConvergenceWarning
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline

# MAX_ITER below is deliberately low to keep `fit` quick. We are timing, not
# fitting a good model, so the resulting warnings are just noise here.
warnings.filterwarnings("ignore", category=ConvergenceWarning)

# ---------------------------------------------------------------- knobs ----

# Edit this line to choose what to run. Unavailable entries are skipped.
DEVICES = ["numpy", "torch-cpu", "torch-cuda", "torch-mps", "cupy"]

N_SAMPLES_FIT = 200_000
N_SAMPLES_PREDICT = 50_000
N_SAMPLES_SCORE = 50_000

N_FEATURES = 100
N_COMPONENTS = 53
N_CLASSES = 11

# torch MPS has no float64. Change to "float64" if you are not benchmarking MPS.
DTYPE = "float32"

# Number of timed repetitions; the fastest one is reported.
N_REPEAT = 5

MAX_ITER = 50

# ------------------------------------------------------------- backends ----


def available(device):
    """Whether `device` can be used here, and why not if it cannot.

    Importing and probing are both wrapped, because a library can be installed
    and still fail to load (a cupy install without a working CUDA driver, for
    instance).
    """
    if device == "numpy":
        return True, ""

    module = {"cupy": "cupy"}.get(device, "torch")
    if importlib.util.find_spec(module) is None:
        return False, f"{module} is not installed"

    try:
        if device == "cupy":
            import cupy

            if cupy.cuda.runtime.getDeviceCount() == 0:
                return False, "cupy reports no CUDA device"
        elif device == "torch-cpu":
            import torch
        elif device == "torch-cuda":
            import torch

            if not torch.cuda.is_available():
                return False, "torch reports no CUDA device"
        elif device == "torch-mps":
            import torch

            if not torch.backends.mps.is_available():
                return False, "torch reports no MPS device"
        else:
            return False, f"unknown device {device!r}"
    except Exception as exc:
        return False, f"{module} failed to load: {type(exc).__name__}: {exc}"

    return True, ""


def to_device(array, device):
    """Put a NumPy array into the namespace and on the device of `device`."""
    if device == "numpy":
        return array
    if device.startswith("torch"):
        import torch

        return torch.asarray(array, device=device.removeprefix("torch-"))
    if device == "cupy":
        import cupy

        return cupy.asarray(array)
    raise ValueError(f"unknown device {device!r}")


def synchronize(device):
    """Block until queued work on `device` has finished."""
    if device == "torch-cuda":
        import torch

        torch.cuda.synchronize()
    elif device == "torch-mps":
        import torch

        torch.mps.synchronize()
    elif device == "cupy":
        import cupy

        cupy.cuda.Device().synchronize()


def to_numpy(array):
    """Bring an array back to NumPy, for the agreement check."""
    if hasattr(array, "detach"):  # torch
        return array.detach().cpu().numpy()
    if hasattr(array, "get"):  # cupy
        return array.get()
    return np.asarray(array)


# ------------------------------------------------------------- plumbing ----


def arm_config(store_as_numpy):
    """Config for one arm of the comparison.

    `store_as_numpy=False` is today's behavior: fitted attributes stay in the
    namespace and on the device of the training data. `True` is the proposal.
    """
    return config_context(array_api_dispatch=True, store_fitted_as_numpy=store_as_numpy)


def measure(fn, device):
    """Time one call of `fn`, in seconds, along with whatever it returned.

    The synchronize is inside the timed region, so GPU numbers are end-to-end
    latencies rather than the time it took to queue the work.
    """
    synchronize(device)
    start = time.perf_counter()
    value = fn()
    synchronize(device)
    return time.perf_counter() - start, value


def make_data(device, random_state):
    """Fit, predict and score datasets, in the namespace of `device`."""
    n_samples = N_SAMPLES_FIT + N_SAMPLES_PREDICT + N_SAMPLES_SCORE
    X, y = make_classification(
        n_samples=n_samples,
        n_features=N_FEATURES,
        n_informative=max(N_FEATURES // 2, 5),
        n_redundant=0,
        n_classes=N_CLASSES,
        n_clusters_per_class=1,
        random_state=random_state,
    )
    X = X.astype(DTYPE)

    fit_end = N_SAMPLES_FIT
    predict_end = fit_end + N_SAMPLES_PREDICT
    return {
        "fit": (to_device(X[:fit_end], device), to_device(y[:fit_end], device)),
        "predict": (to_device(X[fit_end:predict_end], device),),
        "score": (
            to_device(X[predict_end:], device),
            to_device(y[predict_end:], device),
        ),
    }


def run_round(store_as_numpy, data, device):
    """Time one fit, one predict and one score for one arm.

    Returns the times and the predictions, so that the two arms can be compared
    without an extra untimed `predict` call.
    """
    X_fit, y_fit = data["fit"]
    (X_predict,) = data["predict"]
    X_score, y_score = data["score"]

    def fit():
        pipeline = make_pipeline(
            PCA(n_components=N_COMPONENTS),
            LogisticRegression(max_iter=MAX_ITER),
        )
        return pipeline.fit(X_fit, y_fit)

    times = {}
    with arm_config(store_as_numpy):
        times["fit"], pipeline = measure(fit, device)
        times["predict"], predictions = measure(
            lambda: pipeline.predict(X_predict), device
        )
        times["score"], _ = measure(lambda: pipeline.score(X_score, y_score), device)

    return times, to_numpy(predictions)


def fmt(seconds):
    if seconds < 1e-3:
        return f"{seconds * 1e6:.0f}us"
    if seconds < 1:
        return f"{seconds * 1e3:.1f}ms"
    return f"{seconds:.2f}s"


def run_device(device):
    data = make_data(device, random_state=0)

    # Alternate the arms so that a machine that slows down over the run (other
    # load, thermal throttling) does not systematically favor one of them. The
    # first round doubles as the warmup: because we keep the fastest round, the
    # lazy allocation and kernel compilation it pays for are discarded.
    best = {False: {}, True: {}}
    predictions = {}
    for _ in range(N_REPEAT):
        for store_as_numpy in (False, True):
            times, predictions[store_as_numpy] = run_round(store_as_numpy, data, device)
            for method, value in times.items():
                previous = best[store_as_numpy].get(method, float("inf"))
                best[store_as_numpy][method] = min(previous, value)

    print(f"\n{device}")
    print(f"  {'':8}{'as-is':>10}{'as numpy':>12}{'difference':>13}")
    for method in ("fit", "predict", "score"):
        baseline = best[False][method]
        proposal = best[True][method]
        change = (proposal - baseline) / baseline * 100
        print(f"  {method:8}{fmt(baseline):>10}{fmt(proposal):>12}{change:>+12.1f}%")
    agree = np.array_equal(predictions[False], predictions[True])
    print(f"  predictions agree: {'yes' if agree else 'NO'}")


def main():
    print(
        f"fit {N_SAMPLES_FIT} x {N_FEATURES}, predict {N_SAMPLES_PREDICT}, "
        f"score {N_SAMPLES_SCORE}, {N_CLASSES} classes, "
        f"{N_COMPONENTS} components, {DTYPE}"
    )
    print(f"best of {N_REPEAT}, arms alternated")

    for device in DEVICES:
        usable, reason = available(device)
        if usable:
            run_device(device)
        else:
            print(f"\n{device}\n  skipped: {reason}")


if __name__ == "__main__":
    main()