Python高速化 Numba入門 その3
今回もNumbaのドキュメントを読んで行きます。
Numba — numba 0.15.1 documentation
と思ったんですが、読み進めて行くと以外に紹介する内容が少ないことに気づきました。
シンプルなのは良いことなので、今回はUFuncを紹介して、
次回にGPUについて紹介して入門編は一旦終わりたいと思います。
その後に、Numbaを勉強する動機となったものを中心に応用編を話す予定です。
UFunc
UFuncs — numba 0.15.1 documentation
UFunc(ユニバーサル関数)とは、要するにnumpy.ndarrayに対して要素ごとに演算してくれる関数です。
例えばnumpy.logに配列を食わせると、こんな感じです。
>>> import numpy >>> hoge = numpy.arange(1, 5) >>> hoge array([1, 2, 3, 4]) >>> numpy.log(hoge) array([ 0. , 0.69314718, 1.09861229, 1.38629436])
numpy自体にもUFuncの書き方はあるみたいですが、Numbaだとこれもデコレータ一発で行けるようなので見ていきます。
基本は以下の構文。@vectorizeデコレータにリストで型を指定する。
from numba import vectorize import math import numpy @vectorize(['float64(float64, float64)']) def my_ufunc(x, y): return x+y+math.sqrt(x*math.cos(y)) a = numpy.arange(1.0, 10.0) b = numpy.arange(1.0, 10.0) # Calls compiled version of my_ufunc for each element of a and b print(my_ufunc(a, b))
結果は
[ 2.73505259 nan nan nan 11.1909286 14.40021285 16.29724091 nan nan]
マイナスのルート取ろうてして・・・nanが発生。これが例でいいのか・・・
複数の型を受け取りたい時はデコレータの引数のリストにひたすら足していく。
from numba import vectorize import math import numpy @vectorize(['int32(int32, int32)', 'float64(float64, float64)']) def my_ufunc(x, y): return x+y+math.sqrt(x*math.cos(y)) a = numpy.arange(1.0, 10.0, dtype='f8') b = numpy.arange(1.0, 10.0, dtype='f8') print(my_ufunc(a, b)) a = numpy.arange(1, 10, dtype='i4') b = numpy.arange(1, 10, dtype='i4') print(my_ufunc(a, b))
[ 2.73505259 nan nan nan 11.1909286 14.40021285 16.29724091 nan nan] C:\Python27\Scripts\ipython-script.py:16: RuntimeWarning: invalid value encountered in ufunc [ 2 -2147483648 -2147483648 -2147483648 11 14 16 -2147483648 -2147483648]
@jitみたいに@vectorizeにも型推定ないかなと思い実行。
from numba import vectorize import math import numpy @vectorize def my_ufunc(x, y): return x+y+math.sqrt(x*math.cos(y)) a = numpy.arange(1.0, 10.0) b = numpy.arange(1.0, 10.0) # Calls compiled version of my_ufunc for each element of a and b print(my_ufunc(a, b))
TypeError: wrap() takes exactly 1 argument (2 given)
@vectorizeはしっかり型を指定する必要があるらしい。
既存の関数を後から、UFuncBuilderでUFuncにする方法もある。
from numba.npyufunc.ufuncbuilder import UFuncBuilder from numba import f8, i4 import math import numpy def my_ufunc(x, y): return x+y+math.sqrt(x*math.cos(y)) builder = UFuncBuilder(my_ufunc) builder.add(restype=i4, argtypes=[i4, i4]) builder.add(restype=f8, argtypes=[f8, f8]) compiled_ufunc = builder.build_ufunc() a = numpy.arange(1.0, 10.0, dtype='f8') b = numpy.arange(1.0, 10.0, dtype='f8') print(compiled_ufunc(a, b))
Generalized Ufuncs
こちらは私もよく分かっていませんが、わかっている範囲でご紹介します。
通常のUFuncは要素ごとの処理を行っていましたが、Generalized Ufuncsでは内部の配列ベースで処理を行ってくれるそうです。
C言語でも単に要素のforを回すよりも、配列としてBLASに投げた方が速いので動機としては理解できます。
import math import numpy from numba import guvectorize @guvectorize(['void(int32[:,:], int32[:,:], int32[:,:])', 'void(float64[:,:], float64[:,:], float64[:,:])'], '(x, y),(x, y)->(x, y)') def my_gufunc(a, b, c): for i in range(c.shape[0]): for j in range(c.shape[1]): c[i, j] = a[i, j] + b[i, j] a = numpy.arange(1.0, 10.0, dtype='f8').reshape(3,3) b = numpy.arange(1.0, 10.0, dtype='f8').reshape(3,3) # Calls compiled version of my_gufunc for each row of a and b print(my_gufunc(a, b))
型を指定するところは@vectorizeと同じですが、デコレータの第2引数と、本体の関数に戻り値となる引数が増える様です。
デコレータの引数は内部の配列をどう扱うかを指定しています。詳しい話は以下を参考にして下さい。
Generalized Universal Function API — NumPy v1.9 Manual
ちょっと特殊ですが慣れれば行けそうです。
実際にやってみた
前回の@jitと@guvectorizeはどちらが速いか検証します。
# coding: utf-8 import numpy from numba import double from numba.decorators import jit from numba import guvectorize import time import traceback def pairwise_python(X, D): M = X.shape[0] N = X.shape[1] for i in range(M): for j in range(M): d = 0.0 for k in range(N): tmp = X[i, k] - X[j, k] d += tmp * tmp D[i, j] = numpy.sqrt(d) return D @jit def pairwise_numba(X, D): M = X.shape[0] N = X.shape[1] for i in range(M): for j in range(M): d = 0.0 for k in range(N): tmp = X[i, k] - X[j, k] d += tmp * tmp D[i, j] = numpy.sqrt(d) @jit('f8[:, :](f8[:, :], f8[:, :])', nopython=True) def pairwise_numba2(X, D): M = X.shape[0] N = X.shape[1] for i in range(M): for j in range(M): d = 0.0 for k in range(N): tmp = X[i, k] - X[j, k] d += tmp * tmp D[i, j] = numpy.sqrt(d) return D @guvectorize(['void(f8[:, :], f8[:, :])'], '(x, y)->(x, x)') def pairwise_numba3(X, D): M = X.shape[0] N = X.shape[1] for i in range(M): for j in range(M): d = 0.0 for k in range(N): tmp = X[i, k] - X[j, k] d += tmp * tmp D[i, j] = numpy.sqrt(d) if __name__ == '__main__': t = time.time() X = numpy.random.random((1000, 3)) D = numpy.empty((1000, 1000)) pairwise_python(X, D) print "python :", time.time() - t t = time.time() X = numpy.random.random((1000, 3)) D = numpy.empty((1000, 1000)) pairwise_numba(X, D) print "numba_jit :", time.time() - t t = time.time() X = numpy.random.random((1000, 3)) D = numpy.empty((1000, 1000)) pairwise_numba2(X, D) print "numba_jit_type :", time.time() - t t = time.time() X = numpy.random.random((1000, 3)) D = numpy.empty((1000, 1000)) pairwise_numba3(X) print "numba_guvectorize:", time.time() - t
結果は・・・
python : 4.63400006294 numba_jit : 0.131000041962 numba_jit_type : 0.010999917984 numba_guvectorize: 0.0119998455048
特に@guvectorizeだからと言って速いわけではないようだ、
私の書き方が正しくない気もするが・・・