Python: UMAP を使ってみる

UMAP (Uniform Manifold Approximation and Projection) は次元削減手法のひとつ。似た手法としては t-SNE (t-distributed Stochastic Neighbor Embedding) があるけど、それよりも高速らしい。公式のベンチマークが以下で紹介されていて、t-SNE に比べるとスケーラビリティに優れていることが伺える。

umap-learn.readthedocs.io

使った環境は次のとおり。

$ sw_vers 
ProductName:    Mac OS X
ProductVersion: 10.14.6
BuildVersion:   18G5033
$ python -V                          
Python 3.7.7
$ pip list | grep -i umap            
umap-learn                0.4.4

下準備

$ pip install umap-learn scikit-learn

Digit データセットを次元削減してみる

scikit-learn に付属している数字の画像データセットを使って動作を確認してみよう。 UMAP は scikit-learn に準拠したインターフェースを提供している。以下のサンプルコードでは、元が 8 x 8 ピクセルから成る 64 次元の画像データを、UMAP で 2 次元にまで削減している。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
from matplotlib import pyplot as plt
from sklearn import datasets
import umap


def main():
    # データセットを読み込む
    dataset = datasets.load_digits()
    X, y = dataset.data, dataset.target

    # 次元削減する
    mapper = umap.UMAP(random_state=0)
    embedding = mapper.fit_transform(X)

    # 結果を二次元でプロットする
    embedding_x = embedding[:, 0]
    embedding_y = embedding[:, 1]
    for n in np.unique(y):
        plt.scatter(embedding_x[y == n],
                    embedding_y[y == n],
                    label=n)

    # グラフを表示する
    plt.grid()
    plt.legend()
    plt.show()


if __name__ == '__main__':
    main()

上記を保存して実行する。

$ python digitsumap.py

すると、次のようなグラフが得られる。

t-SNE をはじめて使ったときも思ったけど、教師なしでここまで同じラベルのデータがまとまるのは不思議。

t-SNE も試してみる

一応、比較対象として t-SNE も試しておく。以下のサンプルコードは、基本的に先ほどのコードで umap.UMAP() を sklearn.manifold.TSNE に変えただけ。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
from matplotlib import pyplot as plt
from sklearn import datasets
from sklearn import manifold


def main():
    # データセットを読み込む
    dataset = datasets.load_digits()
    X, y = dataset.data, dataset.target

    # 次元削減する
    mapper = manifold.TSNE(random_state=0)
    embedding = mapper.fit_transform(X)

    # 結果を二次元でプロットする
    embedding_x = embedding[:, 0]
    embedding_y = embedding[:, 1]
    for n in np.unique(y):
        plt.scatter(embedding_x[y == n],
                    embedding_y[y == n],
                    label=n)

    # グラフを表示する
    plt.grid()
    plt.legend()
    plt.show()


if __name__ == '__main__':
    main()