矩阵乘法的最佳 numba 实现很大程度上取决于矩阵大小

Question

这个问题与我不久前发布的一个问题有关：
Python, numpy, einsum multiply a stack of matrices

我试图理解为什么在将一堆矩阵相乘时以特定方式使用 Numba 时会得到加速。和以前一样，我放入一个 (500,201,2,2) 数组，沿第一个轴将最后的 (2x2) 矩阵相乘（所以 500 次乘法），得到一个 (201,2,2) 数组作为结果.

这是 Python 代码：

from numba import jit  # numba 0.24, numpy 1.9.3, python 2.7.11

Arr = rand(500,201,2,2)

def loopMult(Arr):
    ArrMult = Arr[0]
    for i in range(1,len(Arr)):
        ArrMult = np.einsum('fij,fjk->fik', ArrMult, Arr[i])
    return ArrMult

@jit(nopython=True)
def loopMultJit(Arr):
    ArrMult = np.empty(shape=Arr.shape[1:], dtype=Arr.dtype)
    for i in range(0, Arr.shape[1]):
        ArrMult[i] = Arr[0, i]
        for j in range(1, Arr.shape[0]):
            ArrMult[i] = np.dot(ArrMult[i], Arr[j, i])
    return ArrMult

@jit(nopython=True)
def loopMultJit_2X2(Arr):
    ArrMult = np.empty(shape=Arr.shape[1:], dtype=Arr.dtype)
    for i in range(0, Arr.shape[1]):
        ArrMult[i] = Arr[0, i]
        for j in range(1, Arr.shape[0]):
            x1 = ArrMult[i,0,0] * Arr[j,i,0,0] + ArrMult[i,0,1] * Arr[j,i,1,0]
            y1 = ArrMult[i,0,0] * Arr[j,i,0,1] + ArrMult[i,0,1] * Arr[j,i,1,1]
            x2 = ArrMult[i,1,0] * Arr[j,i,0,0] + ArrMult[i,1,1] * Arr[j,i,1,0]
            y2 = ArrMult[i,1,0] * Arr[j,i,0,1] + ArrMult[i,1,1] * Arr[j,i,1,1]
            ArrMult[i,0,0] = x1
            ArrMult[i,0,1] = y1
            ArrMult[i,1,0] = x2
            ArrMult[i,1,1] = y2
    return ArrMult

A1 = loopMult(Arr)
A2 = loopMultJit(Arr)
A3 = loopMultJit_2X2(Arr)

print np.allclose(A1, A2)
print np.allclose(A1, A3)

%timeit loopMult(Arr)
%timeit loopMultJit(Arr)
%timeit loopMultJit_2X2(Arr)

这是输出：

True
True
10 loops, best of 3: 40.5 ms per loop
10 loops, best of 3: 36 ms per loop
1000 loops, best of 3: 808 µs per loop

在上一个问题中，接受的答案表明，使用 f2py 可以将速度提高 8 倍，而无需进行详细优化。在这里，使用 Numba，我在 einsum 循环上使用 numba 获得了大约 10% 的加速，但如果不是在循环中使用 np.dot，我只需手动执行 2x2 矩阵乘法，我将获得 45 倍的加速。为什么是这样？我应该提到我已经实现了这两个带有适当类型签名的 jit 函数作为 guvectorize 版本，它基本上提供了相同的加速因子，所以我把它们排除在外。迭代 201,500,2,2 矩阵的加速也很小。

Answer 1

2 评论回应说加速只是由于 python 开销，我认为这是正确的。开销主要是函数调用，还有 for 循环，并且 np.dot 有一些额外的开销。我设置了一个 Naive 点积函数：

@jit(nopython=True)
def dot(mat1, mat2):
    s = 0
    mat = np.empty(shape=(mat1.shape[1], mat2.shape[0]), dtype=mat1.dtype)
    for r1 in range(mat1.shape[0]):
        for c2 in range(mat2.shape[1]):
            s = 0
            for j in range(mat2.shape[0]):
                s += mat1[r1,j] * mat2[j,c2]
            mat[r1,c2] = s
    return mat

然后我设置了函数来乘以数组，一个调用点函数，一个在循环中内置点函数，这样它就可以在没有额外的函数调用的情况下执行：

@jit(nopython=True)
def loopMultJit_dot(Arr):
    ArrMult = np.empty(shape=Arr.shape[1:], dtype=Arr.dtype)
    for i in range(0, Arr.shape[1]):
        ArrMult[i] = Arr[0, i]
        for j in range(1, Arr.shape[0]):
            ArrMult[i] = dot(ArrMult[i], Arr[j, i])
    return ArrMult

@jit(nopython=True)
def loopMultJit_dotInternal(Arr):
    ArrMult = np.empty(shape=Arr.shape[1:], dtype=Arr.dtype)
    for i in range(0, Arr.shape[1]):
        ArrMult[i] = Arr[0, i]
        for j in range(1, Arr.shape[0]):
            s = 0.0
            for r1 in range(ArrMult.shape[1]):
                for c2 in range(Arr.shape[3]):
                    s = 0.0
                    for r2 in range(Arr.shape[2]):
                        s += ArrMult[i,r1,r2] * Arr[j,i,r2,c2]
                    ArrMult[i,r1,c2] = s
    return ArrMult

然后我可以运行 2 个比较：2x2 数组和 10x10 数组。有了这些，我大致了解了为函数调用支付的惩罚，特别是对于 np.dot 函数调用，以及 np.dot 中 BLAS 优化的收益：

print "2x2 Time Test:"
Arr = rand(500,201,2,2)
%timeit loopMult(Arr)
%timeit loopMultJit(Arr)
%timeit loopMultJit_2X2(Arr)
%timeit loopMultJit_dot(Arr)
%timeit loopMultJit_dotInternal(Arr)

print "10x10 Time Test:"
Arr = rand(500,201,10,10)
%timeit loopMult(Arr)
%timeit loopMultJit(Arr)
%timeit loopMultJit_dot(Arr)
%timeit loopMultJit_dotInternal(Arr)

产生：

2x2 Time Test:
10 loops, best of 3: 55.8 ms per loop  # einsum
10 loops, best of 3: 48.7 ms per loop  # np.dot
1000 loops, best of 3: 1.09 ms per loop  # 2x2
10 loops, best of 3: 28.3 ms per loop  # naive dot, separate function
100 loops, best of 3: 2.58 ms per loop  # naive dot internal

10x10 Time Test:
1 loop, best of 3: 499 ms per loop  # einsum
10 loops, best of 3: 91.3 ms per loop  # np.dot
10 loops, best of 3: 170 ms per loop  # naive dot, separate function
10 loops, best of 3: 161 ms per loop  # naive dot internal

我想带回家的信息是：

• 如果您不使用 numba 或需要单行，einsum 很好，但对于矩阵乘法，有更快的选项
• 如果您使用的是小型矩阵，则手动操作会更快，而不是调用单独的函数
• 对于大型矩阵，发明 BLAS 是有原因的，事实上，在小至 10x10 的尺寸下，加速效果非常明显。