めもてふ

Python, Latex, Fortranを取り扱います

数値計算向けpythonの高速化について

Pythonの高速化について

前半では,既に手元に動くコードが有るとして,そのコードの高速化を行う方法を紹介する. 後半では、主にnumba.jitの使い方を説明することになるので AnacondaからPythonを入れておくのが望ましい.

数値計算中のPythonの実行速度が低い原因

数値計算において,Pythonが遅くなる原因は

大きなfor文の中で条件分岐、ランダム値の発生、配列操作を何回も行う

である.しかし,大きなfor文は数値計算する上でどうしても必要なので(特に微分方程式の人) for文の規模はそのままで、高速化する方法について説明することにする.

配列内操作を早くする工夫はnumba.jitを使うのが最も良いであろう。 使い方の詳細は後半で取り上げるが,簡単な例は次節でも説明するので参照されたい.

for文の中の条件分岐・ランダム値の発生に伴う速度低下に対する工夫

このようなfor文の中の条件分岐,ランダム値の発生に対する工夫は2つある.

  1. とりあえずnumba.jitを使う
    • numba.jit とは直下の関数を予めコンパイルする機能.
    • 詳しい使い方は後述する.
  2. 必要そうなら条件分岐を外に出す

これらの工夫を行う例を次に示す.

from numba import jit 
import numpy as np

# iが7の倍数でなく かつ iが5の倍数でない (0 <= i < n)
# jが11の倍数でなく かつ jが6の倍数でない (0 <= j < n) 時に
# i-2*jを足し上げて行く関数

def f1_1(n): # とりあえず組んでみた.遅い
    k = 0
    for i in range(n):
        if i%7 != 0 and i%5 != 0:
            for j in range(n):
                if j%11 != 0 and j%6!= 0:
                    k += i - 2. * j
    return k

@jit
def f1_2(n): # jitをつけてみた.結構早くなった
    k = 0
    for i in range(n):
        if i%7 != 0 and i%5 != 0:
            for j in range(n):
                if j%11 != 0 and j%6!= 0:
                    k += i - 2. * j
    return k

def f1_3(n): # 条件分岐を外に出してみた.もっと早くなった
    k = 0.   # ただしリスト内包表記は@jitの中で使えないみたいなので
             # このように外側に出している.
             # 今回はi_listもj_listも長さが1万以下なので使った.
             # i_listの長さが500万を超えるとちょっとメモリ的にきついイメージがある.
    i_list = np.array([i for i in range(n) if (i%7!=0 and i%5!=0)])
    j_list = np.array([j for j in range(n) if (j%11!=0 and j%6!=0)])
    @jit
    def loop(k):
        for i in i_list:
            for j in j_list:
                k += i - 2. * j
        return k
    return loop(k)
%time print(f1_1(10000))
%time print(f1_2(10000))
%time print(f1_3(10000))

実行結果

-259662261645.0
Wall time: 20.2 s
-259662261645.0
Wall time: 328 ms
-259662261645.0
Wall time: 137 ms
# for文を使う中でどうするかという話をしてるので
# 「次の関数のようにfor文をなくすようにプログラミングせよ」というのは今回はお門違いである.
# (一応一番早そうなのを載せておく)
def f1_5(n):
    k = 0. 
    i_list = np.array([i for i in range(n) if (i%7!=0 and i%5!=0)], dtype=float)
    j_list = np.array([j for j in range(n) if (j%11!=0 and j%6!=0)], dtype=float)
    k= np.sum(i_list)*len(j_list) - 2.* np.sum(j_list)*len(i_list)
    return k
%time print(f1_5(10000))

実行結果

-259662261645.0
Wall time: 9.01 ms

この現象と同様に,for文の中で,1回1回ランダムな値を発生させるのも(jitを使わない時は)遅くなる原因である.

from numba import jit 
import numpy as np

# iが7の倍数でなく かつ iが5の倍数でない (0 <= i < n)
# jが11の倍数でなく かつ jが6の倍数でない (0 <= j < n) 時に
# i-2*j+rを足し上げて行く関数.rはランダムな値

def f2_1(n): # とりあえず組んでみた.遅い
    k = 0
    for i in range(n):
        if i%7 != 0 and i%5 != 0:
            for j in range(n):
                if j%11 != 0 and j%6!= 0:
                    k += i - 2. * j * np.random.rand()
    return k

@jit
def f2_2(n): # jitをつけてみた.結構早くなった
    k = 0
    for i in range(n):
        if i%7 != 0 and i%5 != 0:
            for j in range(n):
                if j%11 != 0 and j%6!= 0:
                    k += i - 2. * j * np.random.rand()
    return k

def f2_3(n):
    k = 0.   
    i_list = np.array([i for i in range(n) if (i%7!=0 and i%5!=0)])
    j_list = np.array([j for j in range(n) if (j%11!=0 and j%6!=0)])
    @jit
    def loop(k):
        j_length = len(j_list)
        for i in i_list: # j のループの前にランダムなベクトルを作成しておく
            r_list = np.random.sample((j_length,))
            for j in j_list:
                k += i - 2. * j * r_list[j]
        return k
    return loop(k)

%time print(f2_1(10000))
%time print(f2_2(10000))
%time print(f2_3(10000))

実行結果

15782474.565609181
Wall time: 26.9 s
51387550.184736826
Wall time: 609 ms
nan
Wall time: 558 ms