More than 5 years have passed since last update.

t-SNE-CUDAで大規模データの超速次元圧縮&可視化

Last updated at 2019-03-16Posted at 2019-03-16

概要

以前t-SNEの説明記事を書きましたが，t-SNE-CUDAなるパッケージが出ていました．
名前の通りt-SNEをCUDAで高速化するものですが，HPML2018(High Performance Machine Learning 2018 Workshop)で発表されました．

単にGPU実装がんばりました，だけではなくアルゴリズム上の工夫もしていて，
SGDの更新式を物理の知見を用いてAttractive ForcesとRepulsive Forcesの項に分けた上で，
以下の各ステップの計算をCUDAで高速化しているようです．

$P_{ij}$の計算
$P_{ij}$と$Q_{ij}$の積
Attractive forcesの計算
Repulsive forcesにおけるBarnes-Hut treeの構築
Replusive forcesのtree走査
低次元空間上の点におけるAttractive/Repulsive forcesの適用

t-SNEはサンプル数が増えると計算時間ものすごく増えるので，高速化されるのはわりと嬉しいですが，
やはり気になるのが「はたしてどのくらい速くなるのか？」
ということで試してみましたのでレポートします．

インストール

詳細は公式リポジトリのInstallationを参照いただければと思いますが，
Condaでのインストールが楽です．とりあえずUbuntu16.04+Anaconda3では成功しました．
ソースからのビルドは試してないですが，依存ライブラリが多く結構めんどくさそうです．

試してみた

とりあえずMNISTでデータ数変えて試してみました．
ただt-SNEと比較するだけだと論文と同じことしてるだけなので，UMAPとの比較もしてみました．

import scipy as sp
from sklearn.datasets import fetch_mldata
import json
from tqdm import tqdm
from tsnecuda import TSNE
import adelheid.plot as pp

import time


if __name__ == "__main__":
    COLORS = ["#000000", "#ffff33", "#66ff66", "#00ccff", "#6633ff", "#cc3399", "#ff3300", "#0066cc", "#9933cc", "#006633"]
    mnist = fetch_mldata("MNIST original", data_home="./")
    n_obs_all = len(mnist["data"])
    n_obs_list = [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 10000, 20000, 30000, 40000, 50000, 60000, n_obs_all]

    result = {}
    for n_obs in tqdm(n_obs_list):
        mnist_data = mnist["data"][:n_obs]
        mnist_labels = mnist["target"][:n_obs].astype(int)

        st = time.clock()

        model = TSNE(n_components=2, perplexity=30.0, theta=0.5, n_iter=1000)
        predicted = model.fit_transform(mnist_data)

        ed = time.clock()
        result[n_obs] = ed-st

        xmin = predicted[:,0].min()
        xmax = predicted[:,0].max()
        ymin = predicted[:,1].min()
        ymax = predicted[:,1].max()

        output_fn = f"mnist_tsnecuda_obs{n_obs}.png"
        fig,ax = pp.subplots(figsize=(16,12))
        unique_labels = sp.unique(mnist_labels)
        for _label in unique_labels:
            inds = sp.where(mnist_labels == _label)[0]
            ax.scatter(predicted[inds,0], predicted[inds,1], c=COLORS[_label], label=_label)

        ax.set_xlim(xmin, xmax)
        ax.set_ylim(ymin, ymax)
        ax.set_xlabel("component 0")
        ax.set_ylabel("component 1")
        ax.set_title("MNIST t-SNE-CUDA visualization")
        ax.legend()
        pp.savefig(output_fn)

    json.dump(result, open("tsnecuda_results.json", "w"), ensure_ascii=False, indent=2, sort_keys=True)

t-SNEやUMAPも含めたテスト用コードはこちらに置いてあります．
(上記コードを動作させるにはこちらのパッケージが必要です)

データ数が10000を超えたあたりからt-SNE-CUDAが有意に速くなりだし，MNISTフルサイズの70000点だと174.46倍ほど高速化されますね．
論文で報告されている数値とほぼ同一です．
UMAPと比較しても4.47倍ほど速いですが，UMAPはt-SNEと比べて各クラスタの中心が離れてクラスタが凝集しやすくなる利点があるので，
このくらいの速度差であれば結局UMAP使ったほうがいいかも，という気もします．

次にt-SNE/t-SNE-CUDA/UMAPで得られた2次元マップを見てみます．