时间戳

首页 > 解决方案 > 如何在 numpy 中使基于序列的函数更快？

python - 如何在 numpy 中使基于序列的函数更快？

问题描述

考虑以下函数：

import numpy as np

a = np.ones(16).reshape(4,4)

def fn(a):
    b = np.array(a)
    for i in range(b.shape[0]):
        for j in range(b.shape[1] - 1):
            b[i][j+1] += b[i][j]
    return b

print(fn(a))

也就是说，对于一个t+1基于t数组计算的通用函数，我可以让它更快吗？我知道有一个np.vectorize但似乎不适合这种情况。

标签： pythonperformancenumpyfor-loopvectorization

解决方案

可以将两个for循环减少为一个for循环，另外只需少量复制开销。

In [86]: a 
Out[86]: 
array([[1., 1., 1., 1.],
       [1., 1., 1., 1.],
       [1., 1., 1., 1.],
       [1., 1., 1., 1.]])

In [87]: b = a.copy() 

In [88]: for col in range(b.shape[1]-1): 
    ...:     b[:, col+1] = np.sum(a[:, :col+2], axis=1) 

In [89]: b
Out[89]: 
array([[1., 2., 3., 4.],
       [1., 2., 3., 4.],
       [1., 2., 3., 4.],
       [1., 2., 3., 4.]])

为了使这个函数适用于通用函数，您可以在 numpy 中寻找等效函数或使用 numpy 操作（矢量化操作）实现一个。对于您提供的示例，我只是使用numpy.sum()它为我们完成了这项工作。

for在性能方面，这种方法比在索引级别使用两个循环操作要好得多，特别是对于较大的数组。在我上面使用的方法中，我们使用列切片。

以下是表明比原生 python 实现速度提高 3 倍以上的时间安排。

原生 Python：

def fn(a):
    b = np.array(a)
    for i in range(b.shape[0]):
        for j in range(b.shape[1] - 1):
            b[i][j+1] += b[i][j]
    return b

稍微矢量化：

In [104]: def slightly_vectorized(b): 
     ...:     for col in range(b.shape[1]-1): 
     ...:         b[:, col+1] = np.sum(a[:, :col+2], axis=1) 
     ...:     return b 

In [100]: a = np.ones(625).reshape(25, 25) 

In [101]: %timeit fn(a) 
303 µs ± 2.05 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [102]: b = a.copy() 

In [103]: %timeit slightly_vectorized(b) 
99.8 µs ± 501 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

推荐阅读