在 numpy 数组上矢量化 python 循环

Question

我需要加快这个循环的处理速度，因为它非常慢。但我不知道如何对其进行矢量化，因为一个值的结果取决于前一个值的结果。有什么建议吗？

import numpy as np

sig = np.random.randn(44100)
alpha = .9887
beta = .999

out = np.zeros_like(sig)

for n in range(1, len(sig)):
    if np.abs(sig[n]) >= out[n-1]:
        out[n] = alpha * out[n-1] + (1 - alpha) * np.abs( sig[n] )
    else:
        out[n] = beta * out[n-1]

Answer 1

Numba 的即时编译器应该通过在首次执行期间将函数编译为本机代码来很好地处理您面临的索引开销：

In [1]: %cpaste
Pasting code; enter '--' alone on the line to stop or use Ctrl-D.
:import numpy as np
:
:sig = np.random.randn(44100)
:alpha = .9887
:beta = .999
:
:def nonvectorized(sig):
:    out = np.zeros_like(sig)
:
:    for n in range(1, len(sig)):
:        if np.abs(sig[n]) >= out[n-1]:
:            out[n] = alpha * out[n-1] + (1 - alpha) * np.abs( sig[n] )
:        else:
:            out[n] = beta * out[n-1]
:    return out
:--

In [2]: nonvectorized(sig)
Out[2]: 
array([ 0.        ,  0.01862503,  0.04124917, ...,  1.2979579 ,
        1.304247  ,  1.30294275])

In [3]: %timeit nonvectorized(sig)
10 loops, best of 3: 80.2 ms per loop

In [4]: from numba import jit

In [5]: vectorized = jit(nonvectorized)

In [6]: np.allclose(vectorized(sig), nonvectorized(sig))
Out[6]: True

In [7]: %timeit vectorized(sig)
1000 loops, best of 3: 249 µs per loop

编辑：按照评论中的建议，添加 jit 基准。 jit(nonvectorized) 正在创建一个轻量级包装器，因此是一种廉价的操作。

In [8]: %timeit jit(nonvectorized)
10000 loops, best of 3: 45.3 µs per loop

函数本身是在第一次执行期间编译的（因此即时）这需要一段时间，但可能不会那么多：

In [9]: %timeit jit(nonvectorized)(sig)
10 loops, best of 3: 169 ms per loop

Answer 2

“前向依赖循环”代码的低向量化潜力

一旦分析了依赖关系，您的“矢量化”并行性中的大部分都将被排除在外。 （JIT 编译器也不能向量化“对抗”这种依赖障碍）

您可以以向量化方式预先计算一些重用值，但没有直接的 Python 语法方式（没有外部 JIT 编译器解决方法）将前移依赖循环计算安排到您的 CPU 向量寄存器中对齐的共并行计算：

from zmq import Stopwatch    # ok to use pyzmq 2.11 for [usec] .Stopwatch()
aStopWATCH =    Stopwatch()  # a performance measurement .Stopwatch() instance

sig    = np.abs(sig)         # self-destructive calc/assign avoids memalloc-OPs
aConst = ( 1 - alpha )       # avoids many repetitive SUB(s) in the loop

for thisPtr in range( 1, len( sig ) ): # FORWARD-SHIFTING-DEPENDENCE LOOP:
    prevPtr = thisPtr - 1              # prevPtr->"previous" TimeSlice in out[] ( re-used 2 x len(sig) times )
    if sig[thisPtr] < out[prevPtr]:                                    # 1st re-use
       out[thisPtr] = out[prevPtr] * beta                              # 2nd
    else:
       out[thisPtr] = out[prevPtr] * alpha + ( aConst * sig[thisPtr] ) # 2nd

在计算策略可以沿原生 numpy 数组的 1D、2D 甚至 3D 结构并行化/广播的情况下，可以看到向量化加速的一个很好的例子。对于大约 100 倍的加速，请参阅Vectorised code for a PNG picture processing ( an OpenGL shader pipeline) 中的 RGBA-2D 矩阵加速处理

性能仍提高了约 3 倍

即使是这个简单的python 代码修订也将速度提高了大约 2.8 倍（现在，即无需进行安装以允许使用临时 JIT 优化编译器）：

>>> def aForwardShiftingDependenceLOOP(): # proposed code-revision
...     aStopWATCH.start()                # ||||||||||||||||||.start
...     for thisPtr in range( 1, len( sig ) ):
...         #        |vvvvvvv|------------# FORWARD-SHIFTING-LOOP DEPENDENCE
...         prevPtr = thisPtr - 1  #|vvvvvvv|--STEP-SHIFTING avoids Numpy syntax
...         if ( sig[ thisPtr] < out[prevPtr] ):
...             out[  thisPtr] = out[prevPtr] * beta
...         else:
...             out[  thisPtr] = out[prevPtr] * alpha + ( aConst * sig[thisPtr] )
...     usec = aStopWATCH.stop()          # ||||||||||||||||||.stop
...     print usec, " [usec]"

>>> aForwardShiftingDependenceLOOP()
57593  [usec]
57879  [usec]
58085  [usec]

>>> def anOriginalForLOOP():
...     aStopWATCH.start()
...     for n in range( 1, len( sig ) ):
...         if ( np.abs( sig[n] ) >= out[n-1] ):
...             out[n] = out[n-1] * alpha + ( 1 - alpha ) * np.abs( sig[n] )
...         else:
...             out[n] = out[n-1] * beta
...     usec = aStopWATCH.stop()
...     print usec, " [usec]"

>>> anOriginalForLOOP()
164907  [usec]
165674  [usec]
165154  [usec]