在numpy数组上矢量化python循环

时间:2014-09-30 18:06:18

标签: python numpy

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

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]

2 个答案:

答案 0 :(得分:3)

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

答案 1 :(得分:1)

"前向相关环路上的低矢量化潜力"代码

你的大部分"矢量化"一旦依赖性被分析,并行性就会脱离游戏。 (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倍的速度(现在,即没有进行安装以允许使用ad-hoc 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]