Python 使用 cython 加速 numpy 矩阵乘法

发布于05月31日

在我的算法中，我计算了几千次矩阵乘法.因此，我计算:

import numpy as np
import time


def mat_mul(mat1, mat2, mat3, mat4):
    return(np.dot(np.transpose(mat1),np.multiply(np.diag(mat2)[:,None], mat3))+mat4)

n = 2000
mat1 = np.random.rand(n,n)
mat2 = np.diag(np.random.rand(n))
mat3 = np.random.rand(n,n)
mat4 = np.random.rand(n,n)

t0=time.time()
cov_11=mat_mul(mat1, mat2, mat1, mat4)
t1=time.time()
print('time ',t1-t0, 's')

矩阵大小:

在我的机器上，我得到以下信息:

我怎样才能加快速度？

谢谢

推荐答案

Numpy实现可以通过减少temporary arrays和reuse them as much as possible的数量(即多次)来稍微优化.实际上，虽然矩阵乘法通常是heavily-optimized by BLAS种实现，但填充/复制(新分配的)数组会增加不可忽略的开销.

以下是实施:

def mat_mul_opt(mat1, mat2, mat3, mat4):
    tmp1 = np.empty((n,n))
    tmp2 = np.empty((n,n))
    vect = np.diag(mat2)[:,None]
    np.dot(np.transpose(mat1),np.multiply(vect, mat3, out=tmp1), out=tmp2)
    np.add(mat4, tmp2, out=tmp1)
    return tmp1

如果可以修改输入矩阵，或者可以在函数外部预先分配tmp1和tmp2(然后多次重用)，则可以进一步优化代码.以下是一个示例:

def mat_mul_opt2(mat1, mat2, mat3, mat4, tmp1, tmp2):
    vect = np.diag(mat2)[:,None]
    np.dot(np.transpose(mat1),np.multiply(vect, mat3, out=tmp1), out=tmp2)
    np.add(mat4, tmp2, out=tmp1)
    return tmp1

以下是我的i5-9600KF处理器(6核)的性能结果:

mat_mul:                 103.6 ms
mat_mul_opt1:             96.7 ms
mat_mul_opt2:             83.5 ms
np.dot time only:         74.4 ms   (kind of practical lower-bound)
Optimal lower bound:      55   ms   (quite optimistic)