为什么 B = numpy.dot(A,x) 通过执行 B[i,:,:] = numpy.dot(A[i,:,:],x) 循环要慢得多？

Question

我得到了一些我无法解释的效率测试结果。

我想组装一个矩阵 B，它的第 i 个条目 B[i,:,:] = A[i,:,:].dot(x)，其中每个 A[i,:,:] 是二维矩阵，x也是。

我可以通过三种方式来测试性能，我制作了随机 (numpy.random.randn) 矩阵 A = (10,1000,1000)，x = (1000,1200)。我得到以下时间结果：

(1) 单个多维点积

B = A.dot(x)

total time: 102.361 s

(2) 循环遍历 i 并执行 2D 点积

   # initialize B = np.zeros([dim1, dim2, dim3])
   for i in range(A.shape[0]):
       B[i,:,:] = A[i,:,:].dot(x)

total time: 0.826 s

(3) numpy.einsum

B3 = np.einsum("ijk, kl -> ijl", A, x)

total time: 8.289 s

所以，选项 (2) 是迄今为止最快的。但是，仅考虑（1）和（2），我看不出它们之间有什么大的区别。循环和做 2D 点积如何能快 124 倍？他们都使用 numpy.dot。有什么见解吗？

我在下面包含了用于上述结果的代码：

import numpy as np
import numpy.random as npr
import time

dim1, dim2, dim3 = 10, 1000, 1200
A = npr.randn(dim1, dim2, dim2)
x = npr.randn(dim2, dim3)

# consider three ways of assembling the same matrix B: B1, B2, B3

t = time.time()
B1 = np.dot(A,x)
td1 = time.time() - t
print "a single dot product of A [shape = (%d, %d, %d)] with x [shape = (%d, %d)] completes in %.3f s" \
  % (A.shape[0], A.shape[1], A.shape[2], x.shape[0], x.shape[1], td1)


B2 = np.zeros([A.shape[0], x.shape[0], x.shape[1]])
t = time.time()
for i in range(A.shape[0]):
    B2[i,:,:] = np.dot(A[i,:,:], x)
td2 = time.time() - t
print "taking %d dot products of 2D dot products A[i,:,:] [shape = (%d, %d)] with x [shape = (%d, %d)] completes in %.3f s" \
  % (A.shape[0], A.shape[1], A.shape[2], x.shape[0], x.shape[1], td2)

t = time.time()
B3 = np.einsum("ijk, kl -> ijl", A, x)
td3 = time.time() - t
print "using np.einsum, it completes in %.3f s" % td3

Answer 1

使用较小的暗角10,100,200，我得到了相似的排名

In [355]: %%timeit
   .....: B=np.zeros((N,M,L))
   .....: for i in range(N):
              B[i,:,:]=np.dot(A[i,:,:],x)
   .....: 
10 loops, best of 3: 22.5 ms per loop
In [356]: timeit np.dot(A,x)
10 loops, best of 3: 44.2 ms per loop
In [357]: timeit np.einsum('ijk,km->ijm',A,x)
10 loops, best of 3: 29 ms per loop

In [367]: timeit np.dot(A.reshape(-1,M),x).reshape(N,M,L)
10 loops, best of 3: 22.1 ms per loop

In [375]: timeit np.tensordot(A,x,(2,0))
10 loops, best of 3: 22.2 ms per loop

迭代速度更快，尽管没有你的情况那么快。

只要迭代维度与其他维度相比较小，这可能是正确的。在这种情况下，与计算时间相比，迭代（函数调用等）的开销很小。一次执行所有值会占用更多内存。

我尝试了dot 的变体，我将A 重塑为二维，认为dot 在内部进行了这种重塑。我很惊讶它实际上是最快的。 tensordot 可能正在进行相同的重塑（如果 Python 可读，则该代码）。

---

einsum 设置涉及 4 个变量的“乘积之和”迭代，i,j,k,m - 即 dim1*dim2*dim2*dim3 与 C 级别 nditer 的步骤。因此，您拥有的索引越多，迭代空间就越大。

Answer 2

numpy.dot 仅代表BLAS 矩阵乘法when the inputs each have dimension at most 2：

#if defined(HAVE_CBLAS)
    if (PyArray_NDIM(ap1) <= 2 && PyArray_NDIM(ap2) <= 2 &&
            (NPY_DOUBLE == typenum || NPY_CDOUBLE == typenum ||
             NPY_FLOAT == typenum || NPY_CFLOAT == typenum)) {
        return cblas_matrixproduct(typenum, ap1, ap2, out);
    }
#endif

当您将整个 3 维 A 数组粘贴到 dot 中时，NumPy 会采用较慢的路径，通过 nditer 对象。 It still tries to get some use out of BLAS 在慢速路径中，但是慢速路径的设计方式，它只能使用向量-向量乘法而不是矩阵-矩阵乘法，这不会给 BLAS 任何接近优化的空间。

Answer 3

我对 numpy 的 C-API 不太熟悉，numpy.dot 就是这样一种内置函数，在早期版本中它曾经位于 _dotblas 之下。

不过，这是我的想法。

1) numpy.dot 对二维数组和 n 维数组采用不同的路径。来自numpy.dot的online documentation：

对于二维数组，它相当于矩阵乘法，对于一维数组，它相当于向量的内积（没有复共轭）。对于 N 维，它是 a 的最后一个轴和 b 的倒数第二个轴的和积

dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m])

因此，对于二维数组，您始终可以保证一次调用 BLAS 的 dgemm，但是对于 ND 数组，numpy 可能会选择可能不对应于最快变化轴的数组的乘法轴（从我发布的摘录），因此可能会错过dgemm 的全部功能。

2) 您的 A 数组太大而无法加载到 CPU 缓存中。在您的示例中，您使用A 和尺寸(10,1000,1000)，这给出了

In [1]: A.nbytes
80000000
In [2]: 80000000/1024
78125

这几乎是80MB，比您的缓存大小大得多。所以你又一次失去了dgemm 的大部分权力。

3) 您还对函数计时有些不精确。众所周知，Python 中的 time 函数并不精确。请改用timeit。

考虑到以上几点，让我们尝试使用可以加载到缓存中的数组

dim1, dim2, dim3 = 20, 20, 20
A = np.random.rand(dim1, dim2, dim2)
x = np.random.rand(dim2, dim3)

def for_dot1(A,x):
    for i in range(A.shape[0]):
        np.dot(A[i,:,:], x)

def for_dot2(A,x):
    for i in range(A.shape[0]):
        np.dot(A[:,i,:], x)    

def for_dot3(A,x):
    for i in range(A.shape[0]):
        np.dot(A[:,:,i], x)  

这是我得到的时间安排（使用 numpy 1.9.2 针对 OpenBLAS 0.2.14 构建的）：

In [3]: %timeit np.dot(A,x)
10000 loops, best of 3: 174 µs per loop
In [4]: %timeit np.einsum("ijk, kl -> ijl", A, x)
10000 loops, best of 3: 108 µs per loop
In [5]: %timeit np.einsum("ijk, lk -> ijl", A, x)
10000 loops, best of 3: 97.1 µs per loop
In [6]: %timeit np.einsum("ikj, kl -> ijl", A, x)
1000 loops, best of 3: 238 µs per loop
In [7]: %timeit np.einsum("kij, kl -> ijl", A, x)
10000 loops, best of 3: 113 µs per loop
In [8]: %timeit for_dot1(A,x)
10000 loops, best of 3: 101 µs per loop
In [9]: %timeit for_dot2(A,x)
10000 loops, best of 3: 131 µs per loop
In [10]: %timeit for_dot3(A,x)
10000 loops, best of 3: 133 µs per loop

请注意，仍然存在时间差，但不是数量级。还要注意choosing the axis of multiplication 的重要性。现在也许，一个 numpy 开发人员可以了解 numpy.dot 在 N-D 阵列的幕后实际上做了什么。