Numpyでfor loopを早くしたい。

Question

Pythonは for文処理が遅く、numpyを利用して早くしたいと思います。

2枚の画像データ(img_kido_dx, img_kido_dy)と
照明の勾配データ(Rpg_grad_p, Rpg_grad_q)を利用し、
各画素の何らかの情報を持つ行列(2x2) Hを計算したいと思います。

アルゴリズム：下記の処理を各画素で行い、各画素に対応する(2x2) Hを求める
matE = [[E_x1,E_x2,...E_xn], [E_y1,E_y2,...E_yn]]
matR = [[R_p1,R_p2,...R_pn], [R_q1,R_q2,...R_qn]]
matM = inv(matR)
matH = matE x matM

普段の場合、下記のfor文で書きますが、やはり遅いです。（約30秒～）
これをnumpyで早くしたい場合、どうすればいいか教えていただけますでしょうか。

import numpy as np

N_ROW = 1080
N_COL = 1920
N_light = 20

img_kido_dx = np.empty((N_ROW * N_COL, N_light)) 
img_kido_dy = np.empty((N_ROW * N_COL,N_light))
Rpg_grad_p = np.empty((N_ROW  * N_COL,N_light))
Rpg_grad_q = np.empty((N_ROW * N_COL,N_light))

for i in range(N_ROW * N_COL):

    mat_E = np.array([img_kido_dx[i, :], img_kido_dy[i, :]])
    mat_R = np.array([Rpq_grad_p[i, :], Rpq_grad_q[i, :]])
    
    mat_M = np.linalg.pinv(mat_R)

    mat_H = np.matmul(mat_E, mat_M)

Answer 1

• データの配置を見直し
• numbaで最適化（データ型の指定およびスレッド並列化）

で10倍弱くらいは高速化できました。

import timeit
import numba as nb
import numpy as np

N_ROW = 200
N_COL = 400
N_light = 20

img_kido_dx = np.empty((N_ROW * N_COL, N_light)) + 0.1
img_kido_dy = np.empty((N_ROW * N_COL,N_light)) + 0.1
Rpg_grad_p = np.empty((N_ROW  * N_COL,N_light)) + 0.1
Rpg_grad_q = np.empty((N_ROW * N_COL,N_light)) + 0.1


def f0():
    mat_H_all = np.empty((N_ROW * N_COL, 2, 2))
    for i in range(N_ROW * N_COL):
        mat_E = np.array([img_kido_dx[i, :], img_kido_dy[i, :]])
        mat_R = np.array([Rpg_grad_p[i, :], Rpg_grad_q[i, :]])
        mat_M = np.linalg.pinv(mat_R)
        mat_H_all[i] = np.matmul(mat_E, mat_M)
    return mat_H_all
     

@nb.jit("(f8[:,:,:], f8[:,:,:], f8[:,:,:])", 
        nopython=True, parallel=True, nogil=True)
def _f1(mat_E_all, mat_R_all, mat_H_all):
    for i in nb.prange(N_ROW * N_COL):
        mat_E = mat_E_all[i]
        mat_R = mat_R_all[i]
        mat_M = np.linalg.pinv(mat_R)
        mat_H_all[i] = np.dot(mat_E, mat_M)

def f1():
    mat_E_all = np.stack([img_kido_dx, img_kido_dy], axis=1)
    mat_R_all = np.stack([Rpg_grad_p, Rpg_grad_q], axis=1)                                
    mat_H_all = np.empty((N_ROW * N_COL, 2, 2))
    _f1(mat_E_all, mat_R_all, mat_H_all)    
    return mat_H_all
    

print(timeit.timeit(f0, number=1))
print(timeit.timeit(f1, number=1))

"""
8.433251421000023
0.9624153400000068
"""

Comments

[kotai2001]

どうもありがとうございました。Numbaで、for文を残しても10倍ほど改善しますね。大変勉強になりました。