python - 在 numpy 中矢量化后的性能损失

Question

我正在编写一个耗时的程序。为了减少时间，我尽力使用 numpy.dot 而不是 for 循环。

但是，我发现矢量化程序的性能比 for 循环版本差得多:

import numpy as np
import datetime
kpt_list = np.zeros((10000,20),dtype='float')
rpt_list = np.zeros((1000,20),dtype='float')
h_r = np.zeros((20,20,1000),dtype='complex')
r_ndegen = np.zeros(1000,dtype='float')
r_ndegen.fill(1)
# setup completed
# this is a the vectorized version
r_ndegen_tile = np.tile(r_ndegen.reshape(1000, 1), 10000)
start = datetime.datetime.now()
phase = np.exp(1j * np.dot(rpt_list, kpt_list.T))/r_ndegen_tile
kpt_data_1 = h_r.dot(phase)
end = datetime.datetime.now()
print((end-start).total_seconds())
# the result is 19.302483
# this is the for loop version
kpt_data_2 = np.zeros((20, 20, 10000), dtype='complex')
start = datetime.datetime.now()
for i in range(10000):
    kpt = kpt_list[i, :]
    phase = np.exp(1j * np.dot(kpt, rpt_list.T))/r_ndegen
    kpt_data_2[:, :, i] = h_r.dot(phase)
end = datetime.datetime.now()
print((end-start).total_seconds())
# the result is 7.74583

这里发生了什么？

Answer 1

我建议您做的第一件事是将您的脚本分解成单独的函数，以便更轻松地进行分析和调试:

def setup(n1=10000, n2=1000, n3=20, seed=None):

    gen = np.random.RandomState(seed)
    kpt_list = gen.randn(n1, n3).astype(np.float)
    rpt_list = gen.randn(n2, n3).astype(np.float)
    h_r = (gen.randn(n3, n3,n2) + 1j*gen.randn(n3, n3,n2)).astype(np.complex)
    r_ndegen = gen.randn(1000).astype(np.float)

    return kpt_list, rpt_list, h_r, r_ndegen


def original_vec(*args, **kwargs):

    kpt_list, rpt_list, h_r, r_ndegen = setup(*args, **kwargs)

    r_ndegen_tile = np.tile(r_ndegen.reshape(1000, 1), 10000)
    phase = np.exp(1j * np.dot(rpt_list, kpt_list.T)) / r_ndegen_tile
    kpt_data = h_r.dot(phase)

    return kpt_data


def original_loop(*args, **kwargs):

    kpt_list, rpt_list, h_r, r_ndegen = setup(*args, **kwargs)

    kpt_data = np.zeros((20, 20, 10000), dtype='complex')
    for i in range(10000):
        kpt = kpt_list[i, :]
        phase = np.exp(1j * np.dot(kpt, rpt_list.T)) / r_ndegen
        kpt_data[:, :, i] = h_r.dot(phase)

    return kpt_data

我还强烈建议使用随机数据而不是全零或全一数组，除非那是您的实际数据的样子(!)。这使得检查代码的正确性变得更加容易 - 例如，如果您的最后一步是乘以一个零矩阵，那么您的输出将始终为全零，无论之前是否存在错误你的代码。

---

接下来，我将通过 line_profiler 运行这些函数查看他们大部分时间都花在哪里。特别是对于 original_vec:

In [1]: %lprun -f original_vec original_vec()
Timer unit: 1e-06 s

Total time: 23.7598 s
File: <ipython-input-24-c57463f84aad>
Function: original_vec at line 12

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    12                                           def original_vec(*args, **kwargs):
    13                                           
    14         1        86498  86498.0      0.4      kpt_list, rpt_list, h_r, r_ndegen = setup(*args, **kwargs)
    15                                           
    16         1        69700  69700.0      0.3      r_ndegen_tile = np.tile(r_ndegen.reshape(1000, 1), 10000)
    17         1      1331947 1331947.0      5.6      phase = np.exp(1j * np.dot(rpt_list, kpt_list.T)) / r_ndegen_tile
    18         1     22271637 22271637.0     93.7      kpt_data = h_r.dot(phase)
    19                                           
    20         1            4      4.0      0.0      return kpt_data

您可以看到它花费了 93% 的时间来计算 h_r 和 phase 之间的点积。这里，h_r 是一个(20, 20, 1000) 数组，phase 是(1000, 10000)。我们正在计算 h_r 的最后一个维度和 phase 的第一个维度的和积(您可以用 einsum 表示法将其写为 ijk,kl->ijl).

---

h_r 的前两个维度在这里并不重要 - 我们可以轻松地将 h_r reshape 为 (20*20, 1000) 点积之前的数组。事实证明，这种 reshape 操作本身可以带来巨大的性能提升:

In [2]: %timeit h_r.dot(phase)
1 loop, best of 3: 22.6 s per loop

In [3]: %timeit h_r.reshape(-1, 1000).dot(phase)
1 loop, best of 3: 1.04 s per loop

我不完全确定为什么会这样——我希望 numpy 的 dot 函数足够智能，可以自动应用这个简单的优化。在我的笔记本电脑上，第二种情况似乎使用了多线程，而第一种情况则没有，这表明它可能没有调用多线程 BLAS 例程。

---

这是一个包含整形操作的矢量化版本:

def new_vec(*args, **kwargs):

    kpt_list, rpt_list, h_r, r_ndegen = setup(*args, **kwargs)

    phase = np.exp(1j * np.dot(rpt_list, kpt_list.T)) / r_ndegen[:, None]
    kpt_data = h_r.reshape(-1, phase.shape[0]).dot(phase)

    return kpt_data.reshape(h_r.shape[:2] + (-1,))

-1 索引告诉 numpy 根据其他维度和数组中元素的数量推断这些维度的大小。我也用过 broadcasting除以 r_ndegen，这样就不需要 np.tile。

通过使用相同的随机输入数据，我们可以检查新版本给出的结果是否与原始版本相同:

In [4]: ans1 = original_loop(seed=0)

In [5]: ans2 = new_vec(seed=0)    
In [6]: np.allclose(ans1, ans2)
Out[6]: True

一些性能基准:

In [7]: %timeit original_loop()
1 loop, best of 3: 13.5 s per loop

In [8]: %timeit original_vec()
1 loop, best of 3: 24.1 s per loop

In [5]: %timeit new_vec()
1 loop, best of 3: 2.49 s per loop

---

更新:

我很好奇为什么 np.dot 对于原始的 (20, 20, 1000) h_r 数组要慢得多，所以我深入研究了 numpy 源代码。 multiarraymodule.c 中实现的逻辑事实证明非常简单:

#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

换句话说，numpy 只是检查输入数组中的任何一个是否有 > 2 个维度，并立即回退到矩阵-矩阵乘法的非 BLAS 实现。检查两个数组的内部维度是否兼容似乎不应该太困难，如果兼容，则将它们视为二维并对它们执行 *gemm 矩阵-矩阵乘法。事实上there's an open feature request for this dating back to 2012 ，如果有任何 numpy 开发人员正在阅读......

与此同时，乘张量时要注意这是一个很好的性能技巧。

---

更新 2:

我忘了 np.tensordot .由于它在二维数组上调用与 np.dot 相同的底层 BLAS 例程，因此它可以获得相同的性能提升，但没有所有那些丑陋的 reshape 操作:

In [6]: %timeit np.tensordot(h_r, phase, axes=1)
1 loop, best of 3: 1.05 s per loop