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为(20*20, 1000)数组.事实证明,这种重塑操作本身可以带来巨大的性能提升:

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的根据其它尺寸和阵列中的元件的数目来推断这些尺寸的大小.我也用广播来划分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实现.看起来检查两个数组的内部尺寸是否兼容似乎不太困难,如果是这样,则将它们视为2D并对它们执行*gemm矩阵 - 矩阵乘法.事实上,如果有任何numpy开发者正在阅读,那么这个可以追溯到2012年的开放式功能请求 ......

与此同时,在倍增张量时,要注意这是一个很好的性能技巧.

更新2:

我忘记了np.tensordot.由于它调用与np.dot2D数组相同的底层BLAS例程,因此它可以实现相同的性能提升,但没有所有那些丑陋的reshape操作:

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