配列計算の処理速度向上検討

Question

質問内容

Pythonで作成した配列計算の処理速度向上を検討しております。
色々試したところ、行きついたのがmultiprocessです。
multiprocessでは処理速度が2倍程度向上しましたが、もう少し速くしたいと考えております。
お手数ですがなにか案があればご回答頂けると幸いです。

試したこと

■ Cython・・・処理速度はほぼ変わりませんでした。
■ numba(jit)・・・処理速度はほぼ変わりませんでした。
■ multiprocess・・・約2倍程度向上。

サンプルコード

import numpy as np
import time as t

start = t.time()
np.random.seed(0)
x = np.random.rand(100)
y = np.random.rand(1000, 1000)
z = np.random.rand(1000, 1000)

a = 0
b = np.zeros((y.shape[0], y.shape[1]))
c = np.zeros((z.shape[0], z.shape[1]))

for j in reversed(range(5)):
    for i in reversed(range(5)):

        p = x[a] * y ** i * z ** j
        q = x[a + 25] * y ** i * z ** j
        b = b + p
        c = c + q
        a = a + 1

print('time : ' + str(round((t.time() - start),5)) + ' [sec]')

上記コードをマルチプロセス化

import numpy as np
import time as t
from multiprocessing import Pool, Array

def f(args):

    x = args[0]
    y = args[1]
    z = args[2]
    a = args[3]
    i = args[4]
    j = args[5]

    p = x[a] * y ** i * z ** j
    q = x[a + 25] * y ** i * z ** j

    return p, q

def main():

    start = t.time()

    p = Pool(processes=8)
    np.random.seed(0)
    x = np.random.rand(100)
    y = np.random.rand(1000, 1000)
    z = np.random.rand(1000, 1000)

    b = np.zeros((y.shape[0], y.shape[1]))
    c = np.zeros((z.shape[0], z.shape[1]))

    ii = []
    jj = []
    aa = np.arange(25)
    for j in reversed(range(5)):
        for i in reversed(range(5)):
            ii.append(i)
            jj.append(j)

    args = [(x, y, z, aa[i], ii[i], jj[i]) for i in range(25)]
    out = p.map(f, args)

    for k in range(25):
        b = b + out[k][0][:]
        c = c + out[k][1][:]

    print('time : ' + str(round((t.time() - start),5)) + ' [sec]')

if __name__ == "__main__":
    main()

Answer 1

python では * よりも ** の方が優先順位が高いため、

p = x[a] * y ** i * z ** j

は

p = x[a] * (y ** i) * (z ** j)

と同じです。 y や z のべき乗を複数回計算するのは無駄ですので、それを予め計算するだけでかなり速くなります。

import numpy as np
import time as t

start = t.time()
np.random.seed(0)
x = np.random.rand(100)
y = np.random.rand(1000, 1000)
z = np.random.rand(1000, 1000)

yi=[y**i for i in range(5)]
zj=[z**j for j in range(5)]

a = 0
b = np.zeros((y.shape[0], y.shape[1]))
c = np.zeros((z.shape[0], z.shape[1]))

for j in reversed(range(5)):
    for i in reversed(range(5)):

        p = x[a] * yi[i] * zj[j]
        q = x[a + 25] * yi[i] * zj[j]
        b = b + p
        c = c + q
        a = a + 1

print('time : ' + str(round((t.time() - start),5)) + ' [sec]')

とすると、手元の環境で 3.81263秒→ 0.70095秒と約5.4倍に速くなりました。

更に、べき乗をそれぞれ別々に求めるのも時間の無駄なので、

yi=[y**i for i in range(5)]
zj=[z**j for j in range(5)]

を

yi=[np.ones_like(y)]
for i in range (1,5):
    yi.append(yi[-1]*y)
zj=[np.ones_like(z)]
for j in range (1,5):
    zj.append(zj[-1]*z)

と修正したところ、0.40384秒にまで速くなりました。更に yi[i] と zj[j] の掛け算も2回行うのはもったいないため、

        p = x[a] * yi[i] * zj[j]
        q = x[a + 25] * yi[i] * zj[j]

を

        yizj=yi[i]*zj[j]
        p = x[a] * yizj
        q = x[a + 25] * yizj

とすると、0.3625秒になりました。ただしこの変更は計算の順番が変わるため、二つの結果の間には若干の計算誤差が有るかもしれません。更に、

        b = b + p
        c = c + q

を

        b += p
        c += q

と書くと、0.33198秒になりました。

Comments

[dream-20xx]

早速のご回答ありがとうございます！
ご提示のコードにて高速化されることを確認できました。
最終的なコードを示すのではなく、順序を追って説明をしてくださり大変感謝いたします。
ご指摘の内容は、高速化という観点で全て知らない内容でしたので、勉強になりました。

なお、ご教授頂いたソースを一部OpenCVとすることによりさらに速くなりました。

yi=[np.ones_like(y)]
for i in range (1,5):
yi.append(yi[-1]*y)
zj=[np.ones_like(z)]
for j in range (1,5):
zj.append(zj[-1]*z)

を

yi=[cv2.pow(y, i) for i in range(5)]
zj=[cv2.pow(z, j) for j in range(5)]