Как ускорить это умножение матрицы Numba - PullRequest
Новая
       9

Как ускорить это умножение матрицы Numba

2 голосов
/ 19 июня 2019

Я пытаюсь воспроизвести матричную факторизацию, используя numba. Вот код: 

import numpy as np
import timeit
from numba import jit, float64, prange


@jit('float64[:,:](float64[:,:],float64[:,:])', parallel=True, nopython=True)
def matmul(A, B):
    C = np.zeros((A.shape[0], B.shape[1]))
    for i in prange(A.shape[0]):
        for j in prange(B.shape[1]):
            for k in range(A.shape[0]):
                C[i,j] = C[i,j] + A[i,k]*B[k,j]
    return C



if __name__ == '__main__':
    m_size = 1000
    num_loops = 10
    A = np.random.rand(m_size, m_size)
    B = np.random.rand(m_size, m_size)

    # Numpy
    start = timeit.default_timer()
    for i in range(num_loops):
        A.dot(B)
    stop = timeit.default_timer()
    execution_time = stop - start
    print("Numpy Executed in ", execution_time)


    # Numba
    start = timeit.default_timer()
    for i in range(num_loops):
        matmul(A, B)
    stop = timeit.default_timer()
    execution_time = stop - start
    print("Numba Executed in ", execution_time)

А вот и вывод: 

Numpy Executed in  0.713342247006949
Numba Executed in  17.631791604988393

В посте, связанном с , выступления numba и numpy были очень близки. Что я делаю неправильно и как я могу улучшить характеристики функции matmul?

...