Kesinの知見置き場

知見を共有していきたいじゃないですか

NumPyとCythonを組み合わせると爆速!

前回の記事の最後にcythonとnumpyを組み合わせても速くならなかったと書いてしまったのですが、@frontier45 さんから公式のチュートリアルをちゃんと読みましょう。と教えていただいたので、自分の勉強がてらブログにも書いておきます。

使用するコードは何のひねりもなく公式のチュートリアルから丸パクリです。


注意:以下のベンチマークMacOS 10.7.3 MacBook Core2 Duo 2.26GHzでPythonC/C++コンパイラはMac標準のPython2.7.1, llvm-gccを使用しています。
Python/Cythonはipythonのtimeitを利用して実行時間を測定しています。
なお、以下の記録はあくまで私の環境、私の実装での記録なので比較の結果は正しいというものではないです。公式のチュートリアルでも実行時間が書かれていますので、そちらの方もご覧下さい

1. コードの変更無し

まずは普通のpythonとただcythonでコンパイルしただけの結果を比較します。前回の結果同様、コードを何も変更しなくても少しは速くなります。

convolve_py.py, convolve_cy.pyx
from __future__ import division
import numpy as np
def naive_convolve(f, g):
    if g.shape[0] % 2 != 1 or g.shape[1] % 2 != 1:
        raise ValueError("Only odd dimensions on filter supported")

    vmax = f.shape[0]
    wmax = f.shape[1]
    smax = g.shape[0]
    tmax = g.shape[1]
    smid = smax // 2
    tmid = tmax // 2
    xmax = vmax + 2*smid
    ymax = wmax + 2*tmid

    h = np.zeros([xmax, ymax], dtype=f.dtype)

    #pythonではxrangeの方が高速なのでここだけ元のコードから変更
    #cythonならばどちらでも変わらないはず
    for x in xrange(xmax):
        for y in xrange(ymax):
            s_from = max(smid - x, -smid)
            s_to = min((xmax - x) - smid, smid + 1)
            t_from = max(tmid - y, -tmid)
            t_to = min((ymax - y) - tmid, tmid + 1)
            value = 0
            for s in xrange(s_from, s_to):
                for t in xrange(t_from, t_to):
                    v = x - smid + s
                    w = y - tmid + t
                    value += g[smid - s, tmid - t] * f[v, w]
            h[x, y] = value
    return h
実行結果
In [1]: import numpy as np
In [2]: N = 100
In [3]: f = np.arange(N*N, dtype=np.int).reshape((N,N))
In [4]: g = np.arange(81, dtype=np.int).reshape((9,9))

In [5]: import convolve_py
In [6]: timeit -n2 -r3 convolve_py.naive_convolve(f,g)
2 loops, best of 3: 2.13 s per loop

In [7]: import convolve_cy
In [8]: timeit -n2 -r3 convolve_cy.naive_convolve(f,g)
2 loops, best of 3: 1.74 s per loop

2. 型情報の追加

次はcdefによって型情報を加えます。

変更する部分を解説していきます。

#numpyをcythonから扱えるようにするためにcimportを使用します
cimport numpy as np

#numpyの配列の型情報を扱いやすくするための代入?
DTYPE = np.int

#Cにおける#defineに相当
ctypedef np.int_t DTYPE_t

numpy配列の型をCから扱うときは通常の型の名前に"_t"を付けるそうで、numpy.intならnumpy.int_t, numpy.int32ならnumpy.int32_tという具合です。このプログラムでは結果を格納するnumpyの配列 hがnumpy.intで、そこに代入するための計算用変数 valueがnumpy.int_tになっています。

#引数の型がnp.ndarrayであることを明示すると呼び出し時に型チェックが行われたり
#f.shape[0]のようなアクセスがより効果的になるらしい
def naive_convolve(np.ndarray f, np.ndarray g)

#変数はなるべく全てcdefにしておく
cdef int vmax = f.shape[0]
…
cdef int x, y, s, t, v, w
...

とくにforループと配列の要素番号に使用するものをcdefしていないとpythonの遅いループのままなので、忘れずにcdefしておきましょう

実行結果
In [9]: import convolve_ndarray
In [10]: timeit -n2 -r3 convolve_ndarray.naive_convolve(f,g)
2 loops, best of 3: 2.68 s per loop

えー!?元のpythonよりも遅いってどういうこと・・・
チュートリアルでの結果と異なってしまい、これについては結局よく分からなかったのですが、ndarray以外の変数だけをcdefにしてやるとちゃんと高速になります

In [11]: import convolve_cdef
In [12]: timeit -n2 -r3 convolve_cdef.naive_convolve(f,g)
2 loops, best of 3: 1.56 s per loop

3. 効率的な配列参照

実はこの段階ではまだnumpy配列のアクセスがボトルネックになっています。これをCレベルにするには配列の型と配列の次元数を教えてやる必要があります。

def naive_convolve(np.ndarray[DTYPE_t, ndim=2] f, np.ndarray[DTYPE_t, ndim=2] g):
…
cdef np.ndarray[DTYPE_t, ndim=2] h = …

このチュートリアルではfとgに2次元のnumpy.int型を渡しているので、引数の型情報にDTYPE_t(defineしてあるので実際はnumpy.int_t)とndim=2を明示する必要があります。どうもバイト数が合っていればいけるみたいで、DTYPE_tの代わりにlongとしても動作しました。が、分かりにくくなるので普通にnumpy配列を宣言した時の型名に"_t"を付けたほうが良さそうです。
numpy.ndarrayに型情報を与えることによって、要素番号にint, unsigned intが使用されていればndarrayへのアクセスがCレベルに高速化されます。

実行結果
In [13]: import convolve_efficient
In [14]: timeit -n2 -r3 convolve_efficient.naive_convolve(f,g)
2 loops, best of 3: 9.02 ms per loop

すげーー!!爆速!!今度はチュートリアル同様に圧倒的に速くなりましたね!(よかった...(^.^; ) 型情報をちゃんと与えることでCython+NumPyの真価が発揮されます。


最終的なコードはこちらです。といっても本家のチュートリアルのコードのまんまですが・・・

from __future__ import division
import numpy as np
cimport numpy as np

DTYPE = np.int
ctypedef np.int_t DTYPE_t

def naive_convolve(np.ndarray[DTYPE_t, ndim=2] f, np.ndarray[DTYPE_t, ndim=2] g):
    if g.shape[0] % 2 != 1 or g.shape[1] % 2 != 1:
        raise ValueError("Only odd dimensions on filter supported")
    assert f.dtype == DTYPE and g.dtype == DTYPE

    cdef int vmax = f.shape[0]
    cdef int wmax = f.shape[1]
    cdef int smax = g.shape[0]
    cdef int tmax = g.shape[1]
    cdef int smid = smax // 2
    cdef int tmid = tmax // 2
    cdef int xmax = vmax + 2*smid
    cdef int ymax = wmax + 2*tmid
    cdef np.ndarray[DTYPE_t, ndim=2] h = np.zeros([xmax, ymax], dtype=DTYPE)
    cdef int x, y, s, t, v, w

    cdef int s_from, s_to, t_from, t_to

    cdef DTYPE_t value
    for x in xrange(xmax):
        for y in xrange(ymax):
            s_from = max(smid - x, -smid)
            s_to = min((xmax - x) - smid, smid + 1)
            t_from = max(tmid - y, -tmid)
            t_to = min((ymax - y) - tmid, tmid + 1)
            value = 0
            for s in xrange(s_from, s_to):
                for t in xrange(t_from, t_to):
                    v = x - smid + s
                    w = y - tmid + t
                    value += g[smid - s, tmid - t] * f[v, w]
            h[x, y] = value
    return h

実行時間比較

time(s) pythonとの相対倍率(倍)
元のPython 2.13 1.0
Cythonコード改変無し 1.74 1.22
型情報追加 2.68 0.79
ndarray以外だけに型情報追加 1.56 1.37
ndarrayに詳細な型情報追加 0.00902 237 (!)


チュートリアルではさらに、@cython.boundscheck(false)のデコレータを使用すると配列参照が配列の範囲内にあるかどうかのチェックを省略するのでさらに高速になると書いてあり、実行時間もさらに半分ぐらいになるようです。
が、私の環境ではチュートリアルのソースを丸コピしても何故か動かせませんでした・・・。
仕方がないので非常に簡単な関数を作って試してみたのですが効果があまり見られず、配列の参照をミスるとC言語で大変お世話になったsegmentation faultの文字と共にpythonを巻き込んで落ちました。ひぃぃ
ちなみに、boundscheck(false)のようなコンパイルオプションは他にもあって高速化に寄与するものもあるみたいです。興味のある方はこちらをどうぞhttp://wiki.cython.org/enhancements/compilerdirectives
自分のスキルではまだ下手にいじれなさそうです^^;

まとめ

  1. ループ変数は確実にcdef intしておく
  2. cimport numpyでCの領域でnumpyを扱う
  3. numpy配列もnp.ndarray[TYPE, ndim=N]の型情報を与える
  4. boundscheck(false)などのコンパイラオプションを与える

ここのチュートリアルでは登場しませんでしたが、他にも

  • 関数をcdefにすることで呼び出しコストを低くする
  • Pythonの関数の代わりにCの関数を呼び出す
  • クラスをcdefにする
  • prangeなどでGILを外して並列化

などの方法で高速化できるらしいです。覚えるのも大変そうですが夢が広がりますね!
それでは皆さんも素敵なCythonライフを!


今回使用したコードをgithubにアップしました