Python 如何加速这个矩阵乘法

我试着用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

在a组中，numba和numpy的表现非常接近。

---

我做错了什么，如何提高matmul函数的性能

为什么不在Numba中简单地调用np.dot（A，B）（实际上是对scipysblas后端的调用）？为更大的矩阵实现一个有效的矩阵乘法并不是那么简单，我的目标是实现一个不同版本的矩阵乘法，在这里，我不求乘积的和，而是求乘积的最小值。此外，我还想对稀疏矩阵这样做。您比较函数性能的帖子使用了一个大小为

（N，3）

的数组

B

，与

（N，N）

（其中

N

较大）的数组相比，它的性能特征非常不同，而且无法利用BLAS在这一制度中使用的算法技巧，在这一制度中，它们发挥了巨大的作用。在不改变你的算法的情况下，我不认为numba可以做任何特别的事情来帮助你。好的，谢谢你，我会尝试另一种方法！