Numpyでcos類似度の計算を高速化する - どこにでもいるSEの備忘録

ただの小ネタです。

個人的にPythonを使っていてcos類似度を計算することがありました。ただ、類似度を計算したいペアの数が多いと結構時間がかかってしまっていました。

これを高速化するにはどうしたらよいか、いろいろ調べて試してみたのでそのメモです。

普通にcos類似度を計算する
- ペア数を巨大にする
高速化パターン
- vectorize
- ndarrayだけを使って計算
使用したノートブック
参考文献
感想

普通にcos類似度を計算する

正直にcos類似度を計算するとこんな感じに書くと思います。

import numpy as np

def cos_sim(v1, v2):
    return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))

X = np.array([0.789, 0.515, 0.335,0])
Y = np.array([0.832, 0.555,0,0])

%time cos_sim(X, Y)

これを実行するとこんな感じになります。

CPU times: user 98 µs, sys: 20 µs, total: 118 µs
Wall time: 121 µs
0.9421693704700748

ペア数を巨大にする

さて、これを下記の条件で実行してみたいと思います。

ベクトルの次元数 : 200
ベクトルの数 : 10_000_000

df = pd.DataFrame({
    "a" : [np.random.random_sample(200)] * 10_000_000,
    "b" : [np.random.random_sample(200)] * 10_000_000
})

%time df.apply(lambda x : cos_sim(x[0], x[1]), axis=1)

※簡略化のため今回はベクトルはすべて同じものを使用しています

これを実行すると出力はこんな感じになります。

CPU times: user 4min 21s, sys: 10.1 s, total: 4min 31s
Wall time: 4min 28s
0          0.738682
1          0.738682
2          0.738682
3          0.738682
4          0.738682
             ...   
9999995    0.738682
9999996    0.738682
9999997    0.738682
9999998    0.738682
9999999    0.738682
Length: 10000000, dtype: float64

10,000,000ペアを計算するのにだいたい4分30秒ほどかかっていることがわかります。実際にはペア数があと1〜2桁多い場合も十分に考えられるので、これ以上かかるのはちょっと許容し難い感じがしますね。

このような、いかにもありがちなケースについて高速化することを考えます。

高速化パターン

vectorize

下記の文献を参考にすると、どうやらnumpyのvectorizeを使うと速くなるらしいのでやってみたいと思います。

python.atelierkobato.com

cos_sim_vectorize = np.vectorize(cos_sim, otypes=[float])

%time cos_sim_vectorize(df["a"], df["b"])

するとこんな感じになります。

CPU times: user 2min 6s, sys: 3.83 s, total: 2min 10s
Wall time: 2min 7s
array([0.73868198, 0.73868198, 0.73868198, ..., 0.73868198, 0.73868198,
       0.73868198])

出力がpd.seriesでないのは置いておくとして、2分10秒ほどで終わってます。当初から考えると、約半分の時間で計算できていることがわかります。

ndarrayだけを使って計算

次に、下記の記事によればnumpyのndarrayの計算だけでも頑張れば書けるそうなので、やってみたいと思います。

www.geeksforgeeks.org

result = np.array([])
chunk = 10 
for i in range(chunk):

    # あんまり行列が巨大だとメモリが不足するので適当に10分割している
    start = int(i * 10_000_000/chunk)
    end = int(i * 10_000_000/chunk + 10_000_000/chunk)

    A = np.array(df.iloc[start : end, 0].values.tolist())
    B = np.array(df.iloc[start : end, 1].values.tolist())

    result = np.concatenate([result, np.sum(A*B, axis=1)/(np.linalg.norm(A, axis=1) * np.linalg.norm(B, axis=1))], 0)

result

ちょっと長くなってしまっていますが、これを実行するとこんな感じです。

CPU times: user 25.4 s, sys: 25.2 s, total: 50.6 s
Wall time: 51.2 s
array([0.73868198, 0.73868198, 0.73868198, ..., 0.73868198, 0.73868198,
       0.73868198])

これで1分切りましたね。当初の素朴な実装から約5~6倍速くなりました。