为什么下面示例中的Python广播比简单循环慢？

Question

我有一个向量数组，并计算它们与第一个向量的差的范数。 使用Python广播时，计算速度比通过简单循环进行计算慢得多。为什么?

import numpy as np

def norm_loop(M, v):
  n = M.shape[0]
  d = np.zeros(n)
  for i in range(n):
    d[i] = np.sum((M[i] - v)**2)
  return d

def norm_bcast(M, v):
  n = M.shape[0]
  d = np.zeros(n)
  d = np.sum((M - v)**2, axis=1)
  return d

M = np.random.random_sample((1000, 10000))
v = M[0]

%timeit norm_loop(M, v) 
25.9 ms

%timeit norm_bcast(M, v)
38.5 ms

我有 Python 3.6.3 和 Numpy 1.14.2。

在 Google Colab 上运行示例的方法如下： https://drive.google.com/file/d/1GKzpLGSqz9eScHYFAuT8wJt4UIZ3ZTru/view?usp=sharing。

- Vladi Raikhlin

Answer 1

内存访问。

首先，广播版本可以简化为

def norm_bcast(M, v):
     return np.sum((M - v)**2, axis=1)

这种方法仍然比循环版本慢一些。 现在，传统智慧认为使用广播的矢量化代码应该总是更快，但在许多情况下并非如此（我会不怕羞地推销我的另一个答案这里）。那么发生了什么？

正如我所说，它涉及内存访问。

在广播版本中，每个M的元素都减去v。处理完最后一行M时，处理第一行的结果已被清除出缓存，所以对于第二步，差异再次被加载到缓存内存并进行平方。最后，它们被加载并处理第三次进行求和。由于M相当大，因此在每个步骤中清除了部分缓存以容纳所有数据。

在循环版本中，每行在一个较小的步骤中完全处理，导致缓存未命中较少且整体代码更快。

最后，可以通过使用einsum的一些数组操作来避免这种情况。此函数允许混合矩阵乘法和求和。首先，我要指出它是一个具有相对于numpy的其余部分而言相当不直观语法的函数，并且潜在的改进通常并不值得为了理解它而付出额外的努力。由于舍入误差，答案可能也略有不同。在这种情况下，可以写成：

def norm_einsum(M, v):
    tmp = M-v
    return np.einsum('ij,ij->i', tmp, tmp)

这将把操作减少到整个数组的两个操作——减法和调用einsum，它执行平方和求和。这样会略微提高效率：

%timeit norm_bcast(M, v)
30.1 ms ± 116 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit norm_loop(M, v)
25.1 ms ± 37.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit norm_einsum(M, v)
21.7 ms ± 65.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

- user2699

---

如果你在norm_loop中改为d[i] = np.linalg.norm(M[i] - v)，你可以节省更多时间。奇怪的是，我用einsum得到的结果比较慢。可能是由于OpenBLAS的原因。 - percusse

Answer 2

最大化性能

在矢量化操作中，您显然具有不良的缓存行为。但是由于未利用现代SIMD指令（AVX2，FMA），计算本身也很慢。幸运的是，克服这些问题并不是非常复杂。

示例

import numpy as np
import numba as nb
@nb.njit(fastmath=True,parallel=True)
def norm_loop_improved(M, v):
  n = M.shape[0]
  d = np.empty(n,dtype=M.dtype)

  #enables SIMD-vectorization 
  #if the arrays are not aligned
  M=np.ascontiguousarray(M)
  v=np.ascontiguousarray(v)

  for i in nb.prange(n):
    dT=0.
    for j in range(v.shape[0]):
      dT+=(M[i,j]-v[j])*(M[i,j]-v[j])
    d[i]=dT
  return d

性能

M = np.random.random_sample((1000, 1000))
norm_loop_improved: 0.11 ms**, 0.28ms
norm_loop: 6.56 ms 
norm_einsum: 3.84 ms

M = np.random.random_sample((10000, 10000))
norm_loop_improved:34 ms
norm_loop: 223 ms
norm_einsum: 379 ms

** 注意性能测量时的细节

第一个结果（0.11ms）是通过反复使用相同数据调用函数得出的。这需要从RAM读取77 GB/s的吞吐量，远超过我的DDR3双通道RAM的能力。由于连续使用相同的输入参数调用函数并不现实，因此我们必须修改测量方法。

为了避免这个问题，我们必须至少两次使用不同的数据调用同一个函数（8MB L3缓存，8MB数据），然后将结果除以2以清除所有缓存。

这些方法的相对性能也因数组大小而异（请查看einsum结果）。

- max9111