tkm2261's blog

研究員(OR屋) → データ分析官 → MLエンジニア → ニート → 米国CS PhDが諸々書いてます

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だからと言って速いわけではないようだ、
私の書き方が正しくない気もするが・・・

次回はGPUの話をする予定
GPU積んでるのがWindows機なので若干手間取るかもしれませんが、頑張ります。