n-th Tensorの距離行列計算を高速でしたい

A8バナー広告

クイックサマリ

  • 手持ちのデータが例えば3rd Tensorだったときに、距離行列を計算する方法を紹介
  • 最大で310倍はやくなった
  • Numpy, Pytorchで試してみた

距離行列？そもそもなんで必要？

ではよくカーネル関数を使う場面に出くわします。古典的にはSVMとか。最近ではGP: Gaussian Processも流行りです。ぼくの同僚たちが研究してるのはGPです。ちょっとウラヤマシイ。もちろん、深層学習もカーネル関数を使うことが「よく」あります。

で、カーネル関数の多くは距離が式の中に入ってます。例えば、ガウスカーネルだと次の式になります。

k(x, y) = exp(- \frac{|| x - y ||^2}{2 \sigma^2})

sigma^2はバンド幅で || x -y ||² が距離ですね。

xとyがスカラ値なら || x – y ||^2 は簡単な話でしょう。でも、もしx, yがベクタだったら？x, yが行列だったら？どうします？

２重ループがもっとも簡単

２重ループで計算しちゃうのが直感的で簡単な実装です。例えば、X, Yを次のように定義しましょう。X, Yはm個のサンプルを持つ集合です。

X = \{x_1,...,x_m \}, Y = \{ y_1,...,y_m \}

距離行列Dはxとyのデータ間距離を意味します。つまり、Dはm行m列の行列です。

D \in \mathbb{R}^{m*m}

じゃあ、２重ループでDを求めるコードを書いてみましょう。

x, yがどんな形状（ベクトルだろうが、行列だろうが）であっても、基本的には２重ループで計算できます。

ボトルネック: Speed

２重ループは簡単ですが、計算速度が問題になります。そもそものループは遅い。

もし、x, yがベクトル、つまり、X, Yが行列の場合（x, yが1-tensorでX, Yが2nd-tensor）は次のようなコードで高速化できます。高速化というか、線形代数的にガウスカーネルの式を置き換えしただけです。

じゃあ、次にx, yが行列の場合を考えてみましょう。つまり、X: (100, 64, 12), Y: (100, 64, 12)です。XとYは100個のサンプルを持つ集合です。１サンプルは(64, 12)の大きさをした行列です。数学的に書くとこうなります。

X = \{x_1,...,x_{100}\}, x_1 \in \mathbb{R}^{64*12}

こうなると、上のコードは使えません。なぜなら np.dot(X, Y.T) っていう転置した二乗計算ができないからです。

で、探しに探すとStackoverflowにエレガントな解法がありました。この解法を使えば、基本的にどんなn-th tensorにでも線形な時間で計算可能です。

アイディアの核になるのは pdict() 関数です。これは距離を高速に計算してくれる実装ですが、入力はベクトルのみです。そこで、データの形を変形し、ベクトルにしてしまおうというわけです。ベクトルになったとしても、indexは一意に保たれますから、情報は可逆です。

1. We reshape (a, b) matrices to a vector having a*b length
2. We compute the distance between vectors.

速度比較

というわけで実装して比較してみました。次の３ケースを比較してます。

  1. with numpy, D(X, X) where X is (50, 64, 12).
  2. with pytorch, D(X, X) where X is (50, 64, 12).
  3. with pytorch, D(X, Y) where X, Y are (50, 64, 12).

コードはGistに載せておきました。どのケースにおいても２重for-loopとreshapeアプローチの解にほとんど差がないことを確認してます。

計算スピードを比較してみましょう。D(X, X)のPytorchケースでは実に350倍も早くなってます。

おしまい

StackOverflowの人。ありがとう(^^ゞ

Python

Posted by user