はじめに
numpy を使った数値計算で、複雑なことをしようとするとつい ndarray ではなく自分で定義した(ndarray を継承していないが例えば包含している)クラスを使いたくなります。しかし、今回の場合はデータのサイズが大きいこともあり、速度面に不安がありました。
そこで、完全に numpy だけで閉じた演算と、(numpy で閉じていないために)リスト内包表記を使った計算でどれくらい計算速度に差が出るのかを調べてみました。
結論を先に書くと、 numpy だけで閉じていたほうが速いです。
追記: ただし、どうしても numpy だけで閉じられないときは、 JIT コンパイルするという手もあるみたいです。
参考: NumPyでfor文を使うと遅いと思ったがそんな事はなかったぜ
コード
なるべくシンプルな状況を考えるために、 [n × 3] のデータを n 個の 3 次元データ点とみなして、それぞれの点のノルムを計算する問題を考えます。
numpy で書くとこんな感じ。
import sys
import numpy as np
n_points = int(sys.argv[1])
data = np.random.random([n_points, 3])
norms = np.linalg.norm(data, axis=1)
axis=1 としているので、全体のノルムではなく(ゼロから数えた) 1 番めの軸に沿ってそれぞれのノルムを計算してくれます。つまり、アウトプットは shape が (n_points,) の ndarray です。
これのリスト内包表記版は以下になります。ひとつの点をひとつの(ndarray やそれを継承していない)オブジェクトである場合に同じ計算をしている想定ですが、なるべく純粋な速度比較をしたいためにその部分は省略しています。
import sys
import numpy as np
n_points = int(sys.argv[1])
data = np.random.random([n_points, 3])
norms = [np.linalg.norm(d) for d in data]
また、この場合の結果は ndarray ではなく長さが n_points の list です。
測定
測定には以下の bash スクリプトを使いました(キレイなコードでなくてすみません)。
# !/bin/bash
name_py=$1
array=(1000 10000 100000 1000000 10000000)
for _ in {1..5}
do
for num in "${array[@]}"
do
{ time python3 "$name_py" "$num" >/dev/null; } 2>&1 | grep real | grep -oE "[0-9]+\.[0-9]+"
done | tr "\n" ","
printf "\n"
done
これをコマンドラインから例えば:
./run_measurements.sh array.py | tee array.txt
などと叩けば、データのサイズをそれぞれ (1000, 10000, 100000, 1000000, 10000000) として先ほどの python コードを実行したものを 5 回やってくれます。
結果
結果を両対数プロットで以下に示します。青が numpy で閉じた版、オレンジがリスト内包表記版で、各点は 5 回の測定の平均です。データのサイズが 10^5 くらいから明らかにリスト内包表記のほうが遅くなっていることが見て取れるかと思います。
また、計測時間の中に乱数生成にかかる時間が含まれているため、スピードの絶対値ではなくあくまでふたつのバージョンの比較のみ意味を持つことに注意してください。
考察
結果はわかったけど、何で!? と思われる方もいらっしゃるかもしれません。そこで、 cProfile という Python のプロファイラを使って中でどういうことが起こっているかみてみましょう。コードはそれぞれ以下のように変更するだけです。
import cProfile
import sys
import numpy as np
n_points = int(sys.argv[1])
data = np.random.random([n_points, 3])
cProfile.run('norms = np.linalg.norm(data, axis=1)')
import cProfile
import sys
import numpy as np
n_points = int(sys.argv[1])
data = np.random.random([n_points, 3])
cProfile.run('norms = [np.linalg.norm(d) for d in data]')
n_points が 100 、 10000 のときの結果をそれぞれ貼ります。
n_points: 100
11 function calls in 0.000 seconds
n_points: 10000
11 function calls in 0.001 seconds
n_points: 100
804 function calls in 0.001 seconds
n_points: 10000
80004 function calls in 0.060 seconds
ご覧のように、 numpy 版ではデータサイズによらず関数呼び出しの回数は一定であるのに対し、リスト内包表記版はデータサイズが増えるに連れて関数呼び出しの回数が増えていっています。リスト内包表記版に関して cProfile の結果を全部載せると以下のようになります(n_points: 100)。
804 function calls in 0.001 seconds
Ordered by: standard name
ncalls tottime percall cumtime percall filename:lineno(function)
1 0.000 0.000 0.001 0.001 <string>:1(<listcomp>)
1 0.000 0.000 0.001 0.001 <string>:1(<module>)
100 0.000 0.000 0.000 0.000 linalg.py:111(isComplexType)
100 0.000 0.000 0.001 0.000 linalg.py:2014(norm)
100 0.000 0.000 0.000 0.000 numeric.py:463(asarray)
1 0.000 0.000 0.001 0.001 {built-in method builtins.exec}
200 0.000 0.000 0.000 0.000 {built-in method builtins.issubclass}
100 0.000 0.000 0.000 0.000 {built-in method numpy.core.multiarray.array}
100 0.000 0.000 0.000 0.000 {built-in method numpy.core.multiarray.dot}
1 0.000 0.000 0.000 0.000 {method 'disable' of '_lsprof.Profiler' objects}
100 0.000 0.000 0.000 0.000 {method 'ravel' of 'numpy.ndarray' objects}
というわけで、 numpy 関連の関数を何度も呼び出していることがわかりますね。つまり、行列計算を得意とする numpy からいちいち python 本体に戻ってまた numpy に行って……というオーバーヘッドを繰り返しているということが見て取れます。
結論
やはり、データサイズが大きい場合は、 numpy の恩恵を最大限に受けるためには計算をなるべく numpy だけで完結させる必要がありそうです。たまに axis
のオプションがない numpy の関数がありますが、なるべくすべてにつけていただきたいところですね。