benchmarks/bench_saga.py - main - githubmirror/scikit-learn

benchmarks/bench_saga.py
390 строк11 KB
Lucy Liu
Deprecate `y_pred` in favour of `y_proba` in `log_loss` (#33740)
16 апр 2026, 11:34
16 апр 2026, 11:34fb4f8fb
"""
Benchmarks of sklearn SAGA vs lightning SAGA vs Liblinear. Shows the gain
in using multinomial logistic regression in term of learning time.
"""

# Authors: The scikit-learn developers
# SPDX-License-Identifier: BSD-3-Clause

import json
import os
import time

import matplotlib.pyplot as plt
import numpy as np

from sklearn.datasets import (
    fetch_20newsgroups_vectorized,
    fetch_rcv1,
    load_digits,
    load_iris,
)
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import log_loss
from sklearn.model_selection import train_test_split
from sklearn.multiclass import OneVsRestClassifier
from sklearn.preprocessing import LabelBinarizer, LabelEncoder
from sklearn.utils.extmath import safe_sparse_dot, softmax
from sklearn.utils.parallel import Parallel, delayed


def fit_single(
    solver,
    X,
    y,
    penalty="l2",
    single_target=True,
    C=1,
    max_iter=10,
    skip_slow=False,
    dtype=np.float64,
):
    if skip_slow and solver == "lightning" and penalty == "l1":
        print("skip_slowping l1 logistic regression with solver lightning.")
        return

    print(
        "Solving %s logistic regression with penalty %s, solver %s."
        % ("binary" if single_target else "multinomial", penalty, solver)
    )

    if solver == "lightning":
        from lightning.classification import SAGAClassifier

    if single_target or solver not in ["sag", "saga"]:
        multi_class = "ovr"
    else:
        multi_class = "multinomial"
    X = X.astype(dtype)
    y = y.astype(dtype)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, random_state=42, stratify=y
    )
    n_samples = X_train.shape[0]
    n_classes = np.unique(y_train).shape[0]
    test_scores = [1]
    train_scores = [1]
    accuracies = [1 / n_classes]
    times = [0]

    if penalty == "l2":
        l1_ratio = 0
        alpha = 1.0 / (C * n_samples)
        beta = 0
        lightning_penalty = None
    else:
        l1_ratio = 1
        alpha = 0.0
        beta = 1.0 / (C * n_samples)
        lightning_penalty = "l1"

    for this_max_iter in range(1, max_iter + 1, 2):
        print(
            "[%s, %s, %s] Max iter: %s"
            % (
                "binary" if single_target else "multinomial",
                penalty,
                solver,
                this_max_iter,
            )
        )
        if solver == "lightning":
            lr = SAGAClassifier(
                loss="log",
                alpha=alpha,
                beta=beta,
                penalty=lightning_penalty,
                tol=-1,
                max_iter=this_max_iter,
            )
        else:
            lr = LogisticRegression(
                solver=solver,
                C=C,
                l1_ratio=l1_ratio,
                fit_intercept=False,
                tol=0,
                max_iter=this_max_iter,
                random_state=42,
            )
            if multi_class == "ovr":
                lr = OneVsRestClassifier(lr)

        # Makes cpu cache even for all fit calls
        X_train.max()
        t0 = time.clock()

        lr.fit(X_train, y_train)
        train_time = time.clock() - t0

        scores = []
        for X, y in [(X_train, y_train), (X_test, y_test)]:
            try:
                y_proba = lr.predict_proba(X)
            except NotImplementedError:
                # Lightning predict_proba is not implemented for n_classes > 2
                y_proba = _predict_proba(lr, X)
            if isinstance(lr, OneVsRestClassifier):
                coef = np.concatenate([est.coef_ for est in lr.estimators_])
            else:
                coef = lr.coef_
            score = log_loss(y, y_proba, normalize=False) / n_samples
            score += 0.5 * alpha * np.sum(coef**2) + beta * np.sum(np.abs(coef))
            scores.append(score)
        train_score, test_score = tuple(scores)

        y_pred = lr.predict(X_test)
        accuracy = np.sum(y_pred == y_test) / y_test.shape[0]
        test_scores.append(test_score)
        train_scores.append(train_score)
        accuracies.append(accuracy)
        times.append(train_time)
    return lr, times, train_scores, test_scores, accuracies


def _predict_proba(lr, X):
    """Predict proba for lightning for n_classes >=3."""
    pred = safe_sparse_dot(X, lr.coef_.T)
    if hasattr(lr, "intercept_"):
        pred += lr.intercept_
    return softmax(pred)


def exp(
    solvers,
    penalty,
    single_target,
    n_samples=30000,
    max_iter=20,
    dataset="rcv1",
    n_jobs=1,
    skip_slow=False,
):
    dtypes_mapping = {
        "float64": np.float64,
        "float32": np.float32,
    }

    if dataset == "rcv1":
        rcv1 = fetch_rcv1()

        lbin = LabelBinarizer()
        lbin.fit(rcv1.target_names)

        X = rcv1.data
        y = rcv1.target
        y = lbin.inverse_transform(y)
        le = LabelEncoder()
        y = le.fit_transform(y)
        if single_target:
            y_n = y.copy()
            y_n[y > 16] = 1
            y_n[y <= 16] = 0
            y = y_n

    elif dataset == "digits":
        X, y = load_digits(return_X_y=True)
        if single_target:
            y_n = y.copy()
            y_n[y < 5] = 1
            y_n[y >= 5] = 0
            y = y_n
    elif dataset == "iris":
        iris = load_iris()
        X, y = iris.data, iris.target
    elif dataset == "20newspaper":
        ng = fetch_20newsgroups_vectorized()
        X = ng.data
        y = ng.target
        if single_target:
            y_n = y.copy()
            y_n[y > 4] = 1
            y_n[y <= 16] = 0
            y = y_n

    X = X[:n_samples]
    y = y[:n_samples]

    out = Parallel(n_jobs=n_jobs, mmap_mode=None)(
        delayed(fit_single)(
            solver,
            X,
            y,
            penalty=penalty,
            single_target=single_target,
            dtype=dtype,
            C=1,
            max_iter=max_iter,
            skip_slow=skip_slow,
        )
        for solver in solvers
        for dtype in dtypes_mapping.values()
    )

    res = []
    idx = 0
    for dtype_name in dtypes_mapping.keys():
        for solver in solvers:
            if not (skip_slow and solver == "lightning" and penalty == "l1"):
                lr, times, train_scores, test_scores, accuracies = out[idx]
                this_res = dict(
                    solver=solver,
                    penalty=penalty,
                    dtype=dtype_name,
                    single_target=single_target,
                    times=times,
                    train_scores=train_scores,
                    test_scores=test_scores,
                    accuracies=accuracies,
                )
                res.append(this_res)
            idx += 1

    with open("bench_saga.json", "w+") as f:
        json.dump(res, f)


def plot(outname=None):
    import pandas as pd

    with open("bench_saga.json", "r") as f:
        f = json.load(f)
    res = pd.DataFrame(f)
    res.set_index(["single_target"], inplace=True)

    grouped = res.groupby(level=["single_target"])

    colors = {"saga": "C0", "liblinear": "C1", "lightning": "C2"}
    linestyles = {"float32": "--", "float64": "-"}
    alpha = {"float64": 0.5, "float32": 1}

    for idx, group in grouped:
        single_target = idx
        fig, axes = plt.subplots(figsize=(12, 4), ncols=4)
        ax = axes[0]

        for scores, times, solver, dtype in zip(
            group["train_scores"], group["times"], group["solver"], group["dtype"]
        ):
            ax.plot(
                times,
                scores,
                label="%s - %s" % (solver, dtype),
                color=colors[solver],
                alpha=alpha[dtype],
                marker=".",
                linestyle=linestyles[dtype],
            )
            ax.axvline(
                times[-1],
                color=colors[solver],
                alpha=alpha[dtype],
                linestyle=linestyles[dtype],
            )
        ax.set_xlabel("Time (s)")
        ax.set_ylabel("Training objective (relative to min)")
        ax.set_yscale("log")

        ax = axes[1]

        for scores, times, solver, dtype in zip(
            group["test_scores"], group["times"], group["solver"], group["dtype"]
        ):
            ax.plot(
                times,
                scores,
                label=solver,
                color=colors[solver],
                linestyle=linestyles[dtype],
                marker=".",
                alpha=alpha[dtype],
            )
            ax.axvline(
                times[-1],
                color=colors[solver],
                alpha=alpha[dtype],
                linestyle=linestyles[dtype],
            )

        ax.set_xlabel("Time (s)")
        ax.set_ylabel("Test objective (relative to min)")
        ax.set_yscale("log")

        ax = axes[2]
        for accuracy, times, solver, dtype in zip(
            group["accuracies"], group["times"], group["solver"], group["dtype"]
        ):
            ax.plot(
                times,
                accuracy,
                label="%s - %s" % (solver, dtype),
                alpha=alpha[dtype],
                marker=".",
                color=colors[solver],
                linestyle=linestyles[dtype],
            )
            ax.axvline(
                times[-1],
                color=colors[solver],
                alpha=alpha[dtype],
                linestyle=linestyles[dtype],
            )

        ax.set_xlabel("Time (s)")
        ax.set_ylabel("Test accuracy")
        ax.legend()
        name = "single_target" if single_target else "multi_target"
        name += "_%s" % penalty
        plt.suptitle(name)
        if outname is None:
            outname = name + ".png"
        fig.tight_layout()
        fig.subplots_adjust(top=0.9)

        ax = axes[3]
        for scores, times, solver, dtype in zip(
            group["train_scores"], group["times"], group["solver"], group["dtype"]
        ):
            ax.plot(
                np.arange(len(scores)),
                scores,
                label="%s - %s" % (solver, dtype),
                marker=".",
                alpha=alpha[dtype],
                color=colors[solver],
                linestyle=linestyles[dtype],
            )

        ax.set_yscale("log")
        ax.set_xlabel("# iterations")
        ax.set_ylabel("Objective function")
        ax.legend()

        plt.savefig(outname)


if __name__ == "__main__":
    solvers = ["saga", "liblinear", "lightning"]
    penalties = ["l1", "l2"]
    n_samples = [100000, 300000, 500000, 800000, None]
    single_target = True
    for penalty in penalties:
        for n_sample in n_samples:
            exp(
                solvers,
                penalty,
                single_target,
                n_samples=n_sample,
                n_jobs=1,
                dataset="rcv1",
                max_iter=10,
            )
            if n_sample is not None:
                outname = "figures/saga_%s_%d.png" % (penalty, n_sample)
            else:
                outname = "figures/saga_%s_all.png" % (penalty,)
            try:
                os.makedirs("figures")
            except OSError:
                pass
            plot(outname)