NumPyでnp.linalg.normを使うよりもnp.einsumを使うほうが速い

3次元空間の2点間の距離は

$$
\sqrt{(x_2 – x_1)^2 + (y_2 – y_1)^2 + (z_2 – z_1)^2}
$$

で計算できますが、2点ではなく多数の点があり、その相互の距離を計算したい場合などは、計算方法によって速度に差が出てきます。それを検証してみます。

比較するソースコード
結果

比較するソースコード

3通りの計算方法を比較します。

NumPyで配列を2乗してsumしてsqrtする
NumPyでnp.linalg.normを使う
NumPyでeinsumしてsqrtする

なお、距離の計算をするのが目的ですがプログラム内では「3成分を2乗して足し合わせてルートする」というノルム計算を行います。

配列を2乗してsumしてsqrtするコード

1億個の3次元ベクトルを用意してノルムを計算します。
まずは数式をそのまま配列計算するように書くパターン。

import numpy as np

try:
    xyz = np.load("xyz.npy")
except:
    N = 100000000
    xyz = np.random.rand(N, 3)
    np.save("xyz.npy", xyz)

norm = np.sqrt(np.sum(xyz**2, axis=1))

NumPyでnp.linalg.normを使うコード

np.linalg.normを使います。ベクトルのノルム計算で調べるとこれが出てくると思います。

import numpy as np

try:
    xyz = np.load("xyz.npy")
except:
    N = 100000000
    xyz = np.random.rand(N, 3)
    np.save("xyz.npy", xyz)

norm = np.linalg.norm(xyz, axis=1)

NumPyでeinsumしてsqrtするコード

このコードでは、行列要素を2乗してsumしてsqrtに入れるという処理を行いますが、書き方がちょっと特殊です。”ij,ij->i”の部分の意味はアインシュタインの縮約記法で検索してみてください。

import numpy as np

try:
    xyz = np.load("xyz.npy")
except:
    N = 100000000
    xyz = np.random.rand(N, 3)
    np.save("xyz.npy", xyz)

norm = np.sqrt(np.einsum("ij,ij->i", xyz, xyz))

結果

以下のような結果になりました。

コード	実行時間（3回平均）
np.sqrt(np.sum(xyz**2, axis=1))	2.30秒
np.linalg.norm(xyz, axis=1)	2.45秒
np.sqrt(np.einsum(“ij,ij->i”, xyz, xyz))	1.43秒

einsumを使うパターンが他よりも2倍くらい速くなりました。一方で、他2つに差がほとんどないというか、むしろnp.linalg.normを使わないほうが若干速いのは意外でした。
距離計算が多数ある場合はこの処理を何度も繰り返すことになるので、そこで高速化ができるのは嬉しいですね。