Примечание

Перейти в конец чтобы скачать полный пример кода или запустить этот пример в браузере через JupyterLite или Binder.

Приближенные ближайшие соседи в TSNE#

Этот пример показывает, как объединить KNeighborsTransformer и TSNE в конвейере. Также демонстрирует, как обернуть пакеты nmslib и pynndescent для замены KNeighborsTransformer и выполнения приближённого поиска ближайших соседей. Эти пакеты можно установить с помощью pip install nmslib pynndescent.

Примечание: в KNeighborsTransformer мы используем определение, которое включает каждую обучающую точку как своего собственного соседа в подсчете n_neighbors, и по причинам совместимости один дополнительный сосед вычисляется, когда mode == 'distance'. Обратите внимание, что мы делаем то же самое в предложенном nmslib обертка.

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

Сначала мы пытаемся импортировать пакеты и предупреждаем пользователя в случае их отсутствия.

import sys

try:
    import nmslib
except ImportError:
    print("The package 'nmslib' is required to run this example.")
    sys.exit()

try:
    from pynndescent import PyNNDescentTransformer
except ImportError:
    print("The package 'pynndescent' is required to run this example.")
    sys.exit()

Мы определяем класс-обертку для реализации API scikit-learn к nmslib, а также функции загрузки.

import joblib
import numpy as np
from scipy.sparse import csr_matrix

from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.datasets import fetch_openml
from sklearn.utils import shuffle


class NMSlibTransformer(TransformerMixin, BaseEstimator):
    """Wrapper for using nmslib as sklearn's KNeighborsTransformer"""

    def __init__(self, n_neighbors=5, metric="euclidean", method="sw-graph", n_jobs=-1):
        self.n_neighbors = n_neighbors
        self.method = method
        self.metric = metric
        self.n_jobs = n_jobs

    def fit(self, X):
        self.n_samples_fit_ = X.shape[0]

        # see more metric in the manual
        # https://github.com/nmslib/nmslib/tree/master/manual
        space = {
            "euclidean": "l2",
            "cosine": "cosinesimil",
            "l1": "l1",
            "l2": "l2",
        }[self.metric]

        self.nmslib_ = nmslib.init(method=self.method, space=space)
        self.nmslib_.addDataPointBatch(X.copy())
        self.nmslib_.createIndex()
        return self

    def transform(self, X):
        n_samples_transform = X.shape[0]

        # For compatibility reasons, as each sample is considered as its own
        # neighbor, one extra neighbor will be computed.
        n_neighbors = self.n_neighbors + 1

        if self.n_jobs < 0:
            # Same handling as done in joblib for negative values of n_jobs:
            # in particular, `n_jobs == -1` means "as many threads as CPUs".
            num_threads = joblib.cpu_count() + self.n_jobs + 1
        else:
            num_threads = self.n_jobs

        results = self.nmslib_.knnQueryBatch(
            X.copy(), k=n_neighbors, num_threads=num_threads
        )
        indices, distances = zip(*results)
        indices, distances = np.vstack(indices), np.vstack(distances)

        indptr = np.arange(0, n_samples_transform * n_neighbors + 1, n_neighbors)
        kneighbors_graph = csr_matrix(
            (distances.ravel(), indices.ravel(), indptr),
            shape=(n_samples_transform, self.n_samples_fit_),
        )

        return kneighbors_graph


def load_mnist(n_samples):
    """Load MNIST, shuffle the data, and return only n_samples."""
    mnist = fetch_openml("mnist_784", as_frame=False)
    X, y = shuffle(mnist.data, mnist.target, random_state=2)
    return X[:n_samples] / 255, y[:n_samples]

Мы сравниваем различные точные/приближенные преобразователи ближайших соседей.

import time

from sklearn.manifold import TSNE
from sklearn.neighbors import KNeighborsTransformer
from sklearn.pipeline import make_pipeline

datasets = [
    ("MNIST_10000", load_mnist(n_samples=10_000)),
    ("MNIST_20000", load_mnist(n_samples=20_000)),
]

max_iter = 500
perplexity = 30
metric = "euclidean"
# TSNE requires a certain number of neighbors which depends on the
# perplexity parameter.
# Add one since we include each sample as its own neighbor.
n_neighbors = int(3.0 * perplexity + 1) + 1

tsne_params = dict(
    init="random",  # pca cannot be used with precomputed distances
    perplexity=perplexity,
    method="barnes_hut",
    random_state=42,
    max_iter=max_iter,
    learning_rate="auto",
)

transformers = [
    (
        "KNeighborsTransformer",
        KNeighborsTransformer(n_neighbors=n_neighbors, mode="distance", metric=metric),
    ),
    (
        "NMSlibTransformer",
        NMSlibTransformer(n_neighbors=n_neighbors, metric=metric),
    ),
    (
        "PyNNDescentTransformer",
        PyNNDescentTransformer(
            n_neighbors=n_neighbors, metric=metric, parallel_batch_queries=True
        ),
    ),
]

for dataset_name, (X, y) in datasets:
    msg = f"Benchmarking on {dataset_name}:"
    print(f"\n{msg}\n" + str("-" * len(msg)))

    for transformer_name, transformer in transformers:
        longest = np.max([len(name) for name, model in transformers])
        start = time.time()
        transformer.fit(X)
        fit_duration = time.time() - start
        print(f"{transformer_name:<{longest}} {fit_duration:.3f} sec (fit)")
        start = time.time()
        Xt = transformer.transform(X)
        transform_duration = time.time() - start
        print(f"{transformer_name:<{longest}} {transform_duration:.3f} sec (transform)")
        if transformer_name == "PyNNDescentTransformer":
            start = time.time()
            Xt = transformer.transform(X)
            transform_duration = time.time() - start
            print(
                f"{transformer_name:<{longest}} {transform_duration:.3f} sec"
                " (transform)"
            )

Пример вывода:

Benchmarking on MNIST_10000:
----------------------------
KNeighborsTransformer  0.007 sec (fit)
KNeighborsTransformer  1.139 sec (transform)
NMSlibTransformer      0.208 sec (fit)
NMSlibTransformer      0.315 sec (transform)
PyNNDescentTransformer 4.823 sec (fit)
PyNNDescentTransformer 4.884 sec (transform)
PyNNDescentTransformer 0.744 sec (transform)

Benchmarking on MNIST_20000:
----------------------------
KNeighborsTransformer  0.011 sec (fit)
KNeighborsTransformer  5.769 sec (transform)
NMSlibTransformer      0.733 sec (fit)
NMSlibTransformer      1.077 sec (transform)
PyNNDescentTransformer 14.448 sec (fit)
PyNNDescentTransformer 7.103 sec (transform)
PyNNDescentTransformer 1.759 sec (transform)

Обратите внимание, что PyNNDescentTransformer занимает больше времени во время первого fit и первый transform из-за накладных расходов компилятора numba just in time. Но после первого вызова скомпилированный код Python сохраняется в кэше numba, и последующие вызовы не страдают от этих начальных накладных расходов. Оба KNeighborsTransformer и NMSlibTransformer выполняются только один раз здесь, так как они показывают более стабильные результаты fit и transform раз (у них нет проблемы холодного старта PyNNDescentTransformer).

import matplotlib.pyplot as plt
from matplotlib.ticker import NullFormatter

transformers = [
    ("TSNE with internal NearestNeighbors", TSNE(metric=metric, **tsne_params)),
    (
        "TSNE with KNeighborsTransformer",
        make_pipeline(
            KNeighborsTransformer(
                n_neighbors=n_neighbors, mode="distance", metric=metric
            ),
            TSNE(metric="precomputed", **tsne_params),
        ),
    ),
    (
        "TSNE with NMSlibTransformer",
        make_pipeline(
            NMSlibTransformer(n_neighbors=n_neighbors, metric=metric),
            TSNE(metric="precomputed", **tsne_params),
        ),
    ),
]

# init the plot
nrows = len(datasets)
ncols = np.sum([1 for name, model in transformers if "TSNE" in name])
fig, axes = plt.subplots(
    nrows=nrows, ncols=ncols, squeeze=False, figsize=(5 * ncols, 4 * nrows)
)
axes = axes.ravel()
i_ax = 0

for dataset_name, (X, y) in datasets:
    msg = f"Benchmarking on {dataset_name}:"
    print(f"\n{msg}\n" + str("-" * len(msg)))

    for transformer_name, transformer in transformers:
        longest = np.max([len(name) for name, model in transformers])
        start = time.time()
        Xt = transformer.fit_transform(X)
        transform_duration = time.time() - start
        print(
            f"{transformer_name:<{longest}} {transform_duration:.3f} sec"
            " (fit_transform)"
        )

        # plot TSNE embedding which should be very similar across methods
        axes[i_ax].set_title(transformer_name + "\non " + dataset_name)
        axes[i_ax].scatter(
            Xt[:, 0],
            Xt[:, 1],
            c=y.astype(np.int32),
            alpha=0.2,
            cmap=plt.cm.viridis,
        )
        axes[i_ax].xaxis.set_major_formatter(NullFormatter())
        axes[i_ax].yaxis.set_major_formatter(NullFormatter())
        axes[i_ax].axis("tight")
        i_ax += 1

fig.tight_layout()
plt.show()

Пример вывода:

Benchmarking on MNIST_10000:
----------------------------
TSNE with internal NearestNeighbors 24.828 sec (fit_transform)
TSNE with KNeighborsTransformer     20.111 sec (fit_transform)
TSNE with NMSlibTransformer         21.757 sec (fit_transform)

Benchmarking on MNIST_20000:
----------------------------
TSNE with internal NearestNeighbors 51.955 sec (fit_transform)
TSNE with KNeighborsTransformer     50.994 sec (fit_transform)
TSNE with NMSlibTransformer         43.536 sec (fit_transform)

Мы можем наблюдать, что по умолчанию TSNE оценщик с его внутренним NearestNeighbors реализация примерно эквивалентна конвейеру с TSNE и KNeighborsTransformer с точки зрения производительности. Это ожидаемо, поскольку оба конвейера внутренне используют один и тот же NearestNeighbors реализация, выполняющая точный поиск соседей. Приближенный NMSlibTransformer уже немного быстрее, чем точный поиск на наименьшем наборе данных, но эта разница в скорости ожидается, что станет более значительной на наборах данных с большим количеством образцов.

Однако обратите внимание, что не все приближённые методы поиска гарантированно улучшают скорость стандартного точного метода поиска: действительно, реализация точного поиска значительно улучшилась с scikit-learn 1.1. Кроме того, метод точного поиска полным перебором не требует построения индекса в fit время. Таким образом, чтобы получить общее улучшение производительности в контексте TSNE pipeline, преимущества приближённого поиска на transform должно быть больше, чем дополнительное время, затраченное на построение приближенного поискового индекса в fit time.

Наконец, сам алгоритм TSNE также вычислительно сложен, независимо от поиска ближайших соседей. Поэтому ускорение поиска ближайших соседей в 5 раз не приведет к ускорению в 5 раз для всего конвейера.