3. 配列計算,NumPy,Numba

3.1 NumPyとNumbaを用いたベクトル和とベクトル内積の計算

NumPyとNumbaを用いてベクトル和とベクトル内積を計算し、 配列計算の基本的な特性を調べます。
リスト3-1にベクトル和とベクトル内積のソースコードを示します。
計算方法として表3-1の3ケースを考えます。

リスト3-1 ベクトル和とベクトル内積のソースコード (vector.py)

"""
vector.py
test program of vector operations (NumPy and Numba)
"""

import time
import numpy as np
from numba import jit, prange, set_num_threads

# (vadd-0) sum of two vectors : for version
def vadd_for(a, b, c):
    n = len(a)
    for i in range(n):
        c[i] = a[i] + b[i]

# (vadd-1) sum of two vectors : array version
def vadd_array(a, b, c):
    c[:] = a + b

# (vadd-2) sum of two vectors : NumPy version
def vadd_numpy(a, b, c):
    c[:] = np.add(a, b)

# (vadd-3) sum of two vectors : Numba version
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def vadd_numba(a, b, c):
    n = len(a)
    for i in prange(n):
        c[i] = a[i] + b[i]

# (sdot-0) scalar product of two vectors : for version
def sdot_for(a, b):
    n = len(a)
    s = 0
    for i in range(n):
        s += a[i] * b[i]
    return s

# (sdot-1) scalar product of two vectors : array version
def sdot_array(a, b):
    return sum(a * b)

# (sdot-2) scalar product of two vectors : NumPy version
def sdot_numpy(a, b):
    return np.dot(a, b)

# (sdot-3) scalar product of two vectors : Numba version
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def sdot_numba(a, b):
    n = len(a)
    s = 0
    for i in prange(n):
        s += a[i] * b[i]
    return s

# parameters
N = 10000000
L = 1000
thread = 8
dtype = 'f4'   # 'f4' or 'f8'
fn = 'vadd-2'
#fn = 'sdot-2'

# Numba threads
if (fn == 'vadd-3') or (fn == 'sdot-3'):
    set_num_threads(thread)

t0 = time.time()

# setup
a = np.arange(N).astype(dtype)
b = np.arange(N).astype(dtype)
c = np.zeros(N, dtype)

t1 = time.time()

# calculation
for _ in range(L):
    if   fn == 'vadd-0':
        vadd_for(a, b, c)
    elif fn == 'vadd-1':
        vadd_array(a, b, c)
    elif fn == 'vadd-2':
        vadd_numpy(a, b, c)
    elif fn == 'vadd-3':
        vadd_numba(a, b, c)
    elif fn == 'sdot-0':
        s = sdot_for(a, b)
    elif fn == 'sdot-1':
        s = sdot_array(a, b)
    elif fn == 'sdot-2':
        s = sdot_numpy(a, b)
    elif fn == 'sdot-3':
        s = sdot_numba(a, b)

t2 = time.time()

# check
if fn.startswith('vadd'):
    r1 = np.sum(c)
    r2 = N * (N - 1)
else:
    r1 = s
    r2 = N * (N - 1) * (2 * N - 1) / 6

# output
print('(%s) N=%d, L=%d' % (fn, N, L))
print('%.2f+%.2f[sec], %e, %e' % (t1 - t0, t2 - t1, r1, r2))

# free
a = None
b = None
c = None

3.2 配列演算について

Pythonでは配列の演算はリスト3-1の関数vadd-1,sdot-1 のようにスカラーのように扱うことができ高速化することができます。 このようなプログラミングを配列指向プログラミングと呼びます。

3.3 NumPyについて

Pythonでは配列の演算はリスト3-1の関数vadd-2,sdot-2のように、 NumPyの関数を使用して計算することができます。
NumPyの関数は高度にチューニングされており、 また自動的にマルチコアを使用して並列処理を行います。

3.4 Numbaについて

Pythonでは前項のような配列演算ができないとき、 for文を用いてそのまま記述すると極めて計算が遅くなります。
そのようなときはNumbaを使用します。
Numbaではリスト3-1の関数vadd-3,sdot3のように関数の前に"@jit"行を挿入するだけで、 JIT(Just In Time)コンパイラーが起動してC言語に変換してコンパイルされ、 C言語と同等の速さになります。
@jit行には"(cache=True, nopython=True, nogil=True, parallel=True)" というオプションを指定することができます。それぞれの意味は以下の通りです。

• cache=True : コンパイルした結果が__pycache__フォルダに保存され、 次回同じ関数を使用するときにそれを使用するので高速に関数を起動することができます。 ただし、関数の引数が変わったときはwarningを出して再コンパイルします。
• nopython=True : コンパイルに失敗したときはエラーメッセージを出して中止します。
• nogil=True : GIL(Global Interpreter Lock)を無効にしてマルチスレッドを有効にします。
• parallel=True : マルチスレッドにより並列計算します。

Numbaを使用する注意点は以下の通りです。 (将来のバージョンでは変わる可能性があります)

• 関数の引数はNumPyのサポートする型（'i4','f4','f8','c8','c16'等）に限定される。 配列を含む構造体など使えない型がある。
• Numbaを使用する関数内で呼び出される関数も@jit宣言されていなければならない。
• for文を並列処理するにはrangeをprangeに変更する。 （import numba.prangeが必要）
• 並列計算のスレッド数を指定するにはnumba.set_num_threads関数を使用する。 スレッド数の既定値は論理コア数。

3.5 計算時間

リスト3-1において、ベクトルの大きさNと繰り返し回数Lを変え、 計算方法にNumPyとNumbaをとり、Numbaについてはスレッド数を変えたときの計算時間を、 表3-2（ベクトル和）と表3-3（ベクトル内積）に示します。
計算時間は計算の本体部であり、配列を用意する部分は除いています。
計算時間の誤差を小さくするために繰り返し回数分だけ計算しています。
ベクトルの大きさNと繰り返し回数Lの積はNo.1～No.4で一定10¹⁰です。 従ってNo.1～No.4の演算量は同じです。
最適化の効果が見られるケースに色をつけています。

(注1)1/100の繰り返し回数時の計算時間を100倍した推定値

表3-2～表3-3から以下のことがわかります。

• 使用メモリーの少ないNo.2～4ではキャッシュメモリーが有効利用され計算時間が大幅に短縮されることがある。
• ベクトル和では配列演算(vadd-1)とNumPy(vadd-2)の計算時間はほぼ同じである。
• ベクトル和(No.1～3)ではNumba(vadd-3)が配列演算(vadd-1)とNumPy(vadd-2)より速い。
• ベクトル内積ではNumPy(sdot-2)と並列化したNumba(sdot-3)はほぼ同じである。
• ベクトル内積の配列演算(sdot-1)は非常に遅い。
• 繰り返し回数が非常に大きいとき(No.4)は余分に時間がかかる。関数を起動するオーバーヘッドの影響と考えられる。

3.6 for文の計算時間

Pythonではfor文を用いてそのまま記述すると極めて計算が遅くなります。
表3-4にfor文のベクトル和(vadd-0)とベクトル内積(sdot-0)の計算時間を示します。
表3-2、表3-3と比べて繰り返し回数Lが1/1000となっていることに注意してください。

表3-4から以下のことがわかります。

• for文では最適化が行われないためにすべての組み合わせの計算時間がほぼ同じである。
• 表3-4の計算時間は表3-3とほぼ同じオーダーなので、 for文の計算速度はNumPyまたはNumbaの1/1000のオーダーである。