目的
Cythonの関数の書き方によって実行速度がどう変わるか調べました。
履歴
- '20/5/8 初出
- '20/5/10 修正(%%cython以下と%timeit以下を別セルで実行することを明示, float64→double, 備考はJupyter環境で動作検証していなかったため動くように修正)
背景
- Pythonで数値計算する場合には、高速化のためにNumpyを使います。Numpyで簡単に書けない場合にはNumbaやCythonを使いますが、Numbaは効果が出ない場合がありますし、Cythonは書き方によって実行時間が大きく変わります。
- Cythonによる高速化は型付けが全てのように思いがちですが、実際にはそれ以外にも注意すべき点があります。それに気づかないと「Pythonは速度は全然ダメ」と早合点して、他の言語を勉強しているうちに時間が経ってしまうということにもなりかねません。
- そこで、書き方による実行時間の違いを、系統的に評価してみました。
方法
- $n$個の実数 $v_{1}, \ldots, v_{n}$ の二乗和 $s=v_{1}^{2}+v_{2}^{2}+\cdots+v_{n}^{2}$ を計算する問題を考えます。単純化のために $v_{1}=v_{2}=\cdots=v_{n}=1.0$、$n=10^6$ とします。
- 実行速度に影響を及ぼす因子として、下の表に示す5種類の要因を考えます。これらの要因を全て網羅した17種類のコードを用意して、実行時間を評価しました1。
- 環境: Python 3.7 / Cython 0.29.17 / Numpy 1.18.1 / Jupyter Notebook
# | 要因 | 選択肢 |
---|---|---|
1 | 使用する構文 | ①Pythonのfor文 ②Cythonのfor文 |
2 | 引数の型 | ①任意のイテラブル ②Numpy配列 ③型付きメモリービュー |
3 | 引数の要素の型指定 | ①なし ②あり |
4 | 要素アクセスの方法 | ①非インデックス指定(v in vs) ②インデックス指定(vs[i]) |
5 | チェック機能の有無 | ①なし ②あり |
結果
- 型指定だけではほとんど高速化しませんでしたが、型指定してインデックスアクセスした場合に10~20倍程度の速度向上が見られました。その他の要因の影響は誤差に毛が生えた程度でした。
- そのため、NumpyやNumbaからCythonに移行する場合には、本当はブロードキャストで書きたかったところを
for i in range(vs.shape[0]):
として、vs[i]
の形のインデックスアクセスをするのが良いようです。リスト内包表記の感覚でfor v in vs:
と書くのは避けるべきです。
詳細
1. 準備
次の数列の二乗和をとります。PythonリストとNumpy配列の2種類のバージョンを用意します。
import numpy as np
vs_list = [1.0,]*10**6
vs_ndarray = np.ones((10**6,), dtype=np.double)
2. Python の for 文
もとになるPythonのforループの性能は以下の通りでした2。$10^6$個の実数値の二乗和をとるのに300~400msかかっています3。Pythonでは、インデックスアクセスせずにイテレータを素直に使うほうが高速なようです。
# | 構文 | 引数の型 | 要素の型指定 | 要素アクセス | チェック機能 | 実行時間[ms] |
---|---|---|---|---|---|---|
1 | Python の for 文 | 任意のイテラブル | なし | v in vs | あり | 302 |
2 | Python の for 文 | 任意のイテラブル | なし | vs[i] | あり | 381 |
def sum_py(vs):
s = 0
for v in vs:
s += v
return s
%timeit square_sum_py(vs_list)
def square_sum_py_range(vs):
s = 0
for i in range(len(vs)):
s += vs[i]**2
return s
%timeit square_sum_py_range(vs_list)
3. Numpyのブロードキャスト
当然のことながら、Numpyのブロードキャストを使うと劇的に高速化できます。しかし、実際にCythonを使うのはブロードキャスト機能が使えない場合なので、ここではそれ以外の方法を考える必要があります。なお、Numpy配列をPython のfor文で処理すると、逆に遅くなってしまうことが分かります。
# | 構文 | 引数の型 | 要素の型指定 | 要素アクセス | チェック機能 | 実行時間[ms] |
---|---|---|---|---|---|---|
3 | Numpyのブロードキャスト | Numpy配列 | あり | なし(不要) | あり | 17 |
4 | Python の for 文 | Numpy配列 | あり | v in vs | あり | 1640 |
5 | Python の for 文 | Numpy配列 | あり | vs[i] | あり | 1950 |
def square_sum_np(vs):
return np.sum(vs**2)
%timeit square_sum_np(vs_ndarray)
# 関数は上で定義
%timeit square_sum_py(vs_ndarray)
# 関数は上で定義
%timeit square_sum_py_range(vs_ndarray)
4. 引数の型指定
そこで、次に、Cythonを使ってPythonのfor文を高速化することを考えます。Cythonで関数に低レベルのメモリアクセスが可能な配列を渡す方法としては、Numpy配列を直接引数に指定する方法と、型付きメモリビューを引数に指定する方法があります。ここでは、より直観的な前者の方法から調べていきます。
ネット上の断片的な情報を見ていると、Cythonで高速化に効くのは型指定だけであるように思えてきます。しかし、下の表を見ると分かるように、関数の引数をNumpy配列に指定しただけでは、元のPythonのコード(1)と比べて全然速くなりません(6)。この段階では、Numpy配列の要素の型を指定してもほとんど改善されません(7)。もちろん、Numpy配列を使ったPythonのコード(4)よりは高速ですが、それだけはCythonを使う理由がありません。
# | 構文 | 引数の型 | 要素の型指定 | 要素アクセス | チェック機能 | 実行時間[ms] |
---|---|---|---|---|---|---|
6 | Cython の for 文 | Numpy配列 | なし | v in vs | あり | 378 |
7 | Cython の for 文 | Numpy配列 | あり | v in vs | あり | 362 |
%%cython
cimport numpy as np
def square_sum_cy_np(np.ndarray vs):
cdef double v, s = 0.0
for v in vs:
s += v**2
return s
%timeit square_sum_cy_np(vs_ndarray)
%%cython
cimport numpy as np
def square_sum_cy_np_typed(np.ndarray[np.double_t, ndim=1] vs):
cdef double v, s = 0.0
for v in vs:
s += v**2
return s
%timeit square_sum_cy_np_typed(vs_ndarray)
5. インデックスアクセス
それでは何をすればよいかというと、配列要素へのアクセス方法(for文の書き方)を変えることです。要素の型指定なしでインデックスアクセスすると遅くなりますが、要素の型指定をしてインデックスアクセスにすると劇的に速くなります。
# | 構文 | 引数の型 | 要素の型指定 | 要素アクセス | チェック機能 | 実行時間[ms] |
---|---|---|---|---|---|---|
8 | Cython の for 文 | Numpy配列 | なし | vs[i] | あり | 1610 |
9 | Cython の for 文 | Numpy配列 | あり | vs[i] | あり | 28 |
%%cython
cimport numpy as np
def square_sum_cy_np_range(np.ndarray vs):
cdef double s = 0.0
for i in range(vs.shape[0]):
s += vs[i]**2
return s
%timeit square_sum_cy_np_range(vs_ndarray)
%%cython
cimport numpy as np
def square_sum_cy_np_typed_range(np.ndarray[np.double_t, ndim=1] vs):
cdef double s = 0.0
for i in range(vs.shape[0]):
s += vs[i]**2
return s
%timeit square_sum_cy_np_typed_range(vs_ndarray)
6. チェック機能の省略
公式ドキュメントには、配列アクセスのチェック機能を省略する方法も載っていますが(10~13)、チェック機能を省略しない場合(6~9)との違いはわずかでした。確かに1~2割は早くなっていますが、最後の一押しという程度の効果しかありません。
# | 構文 | 引数の型 | 要素の型指定 | 要素アクセス | チェック機能 | 実行時間[ms] |
---|---|---|---|---|---|---|
10 | Cython の for 文 | Numpy配列 | なし | v in vs | なし | 315 |
11 | Cython の for 文 | Numpy配列 | あり | v in vs | なし | 313 |
12 | Cython の for 文 | Numpy配列 | なし | vs[i] | なし | 1610 |
13 | Cython の for 文 | Numpy配列 | あり | vs[i] | なし | 25 |
%%cython
cimport numpy as np
from cython import boundscheck, wraparound
def square_sum_cy_np_nocheck(np.ndarray vs):
cdef double v, s = 0.0
with boundscheck(False), wraparound(False):
for v in vs:
s += v**2
return s
%timeit square_sum_cy_np_nocheck(vs_ndarray)
%%cython
cimport numpy as np
from cython import boundscheck, wraparound
def square_sum_cy_np_typed_nocheck(np.ndarray[np.double_t, ndim=1] a):
cdef double d, s = 0.0
with boundscheck(False), wraparound(False):
for d in a:
s += d**2
return s
%timeit square_sum_cy_np_typed_nocheck(vs_ndarray)
%%cython
cimport numpy as np
from cython import boundscheck, wraparound
def square_sum_cy_np_range_nocheck(np.ndarray a):
cdef double s = 0.0
with boundscheck(False), wraparound(False):
for i in range(a.shape[0]):
s += a[i]**2
return s
%timeit square_sum_cy_np_range_nocheck(vs_ndarray)
%%cython
cimport numpy as np
from cython import boundscheck, wraparound
def square_sum_cy_np_typed_range_nocheck(np.ndarray[np.double_t, ndim=1] a):
cdef double s = 0.0
with boundscheck(False), wraparound(False):
for i in range(a.shape[0]):
s += a[i]**2
return s
%timeit square_sum_cy_np_typed_range_nocheck(vs_ndarray)
型付きメモリービュー
Cython関数に低レベルメモリアクセスが可能な配列を渡す方法としては、引数に型付きメモリービューを指定する方法もあります。公式文書ではこちらの方法が推奨されています。以下の結果を見ると分かるように、この場合にも効いているのは配列要素へのアクセスの方法であることが分かります。
# | 構文 | 引数の型 | 要素の型指定 | 要素アクセス | チェック機能 | 実行時間[ms] |
---|---|---|---|---|---|---|
14 | Cython の for 文 | 型付きメモリービュー | あり(必要) | v in vs | あり | 519 |
15 | Cython の for 文 | 型付きメモリービュー | あり(必要) | vs[i] | あり | 26 |
16 | Cython の for 文 | 型付きメモリービュー | あり(必要) | v in vs | なし | 516 |
17 | Cython for 文 | 型付きメモリービュー | あり(必要) | vs[i] | なし | 24 |
%%cython
def square_sum_cy_mv(double[:] vs):
cdef double v, s = 0.0
for v in vs:
s += v**2
return s
%timeit square_sum_cy_np_typed_range_nocheck(vs_ndarray)
%%cython
def square_sum_cy_mv_range(double[:] vs):
cdef double s = 0.0
for i in range(vs.shape[0]):
s += vs[i]**2
return s
%timeit square_sum_cy_mv_range(vs_ndarray)
%%cython
from cython import boundscheck, wraparound
def square_sum_cy_mv_nocheck(double[:] vs):
cdef double v, s = 0.0
with boundscheck(False), wraparound(False):
for v in vs:
s += v**2
return s
%timeit square_sum_cy_mv_nocheck(vs_ndarray)
%%cython
from cython import boundscheck, wraparound
def square_sum_cy_mv_range_nocheck(double[:] vs):
cdef double s = 0.0
with boundscheck(False), wraparound(False):
for i in range(vs.shape[0]):
s += vs[i]**2
return s
%timeit square_sum_cy_mv_range_nocheck(vs_ndarray)
考察
- 以上の結果を見ると、Numpyで書けないところをCythonで補う場合には、15か17のような書き方をするのが良さそうです。もちろん、問題の種類によって傾向は違うはずですが、次にCythonを使う時にはこの辺を最初の一歩にしようかなと思います。
- 懸念事項として、評価に使用したCPUは若干遅すぎたかもしれません3。また、GPUを使って多並列の計算をする場合には最初から考え直す必要があります。
備考
- 公式文書にもあるように、引数として渡す配列(vs)のメモリ配置をC型に指定するとさらに若干高速になります。明示的にコンパイルする方法では気にしていませんでしたが、①配列vsをCython関数の外側で計算する、②Jupyterのマジックコマンドを使う、という制約をつけると、簡単に行かないようです。この辺は整理できていないので、分かったら追加します。
- 多次元配列では次のような書き方をするのがお手軽なようです。
import numpy as np
%load_ext Cython
vs = np.ones((10**3,10**3), dtype=np.double)
%%cython
from cython import boundscheck, wraparound
cdef double square_sum(double[:, :] vs):
# def square_sum(double[:, :] vs): としてもこの場合にはほとんど変わらない
cdef:
double s = 0.0
Py_ssize_t nx = vs.shape[0]
Py_ssize_t ny = vs.shape[1]
Py_ssize_t i, j
with boundscheck(False), wraparound(False):
for i in range(nx):
for j in range(ny):
s += vs[i, j]**2
return s
%timeit square_sum(vs)