加速numpy嵌套循环

时间:2014-05-09 13:14:20

标签: python for-loop numpy slice

我正在使用numpy和cython为python中的无线网络编写模拟,假设在2d平面上随机散布了许多节点no_nodes,这些节点发出一些波形及其各自的接收器,再次分散随机在2d飞机上。每个发送节点产生一个我称之为output的波形(每个波形都可以产生不同长度的输出)。

我想要做的是将每个节点的输出相加到一个大波形,这个波形将成为每个接收器的输入以进行解调等。现在有两个关键点:

  • 发送器异步发送,因此每个发送节点必须维护start_clockend_clock,以便正确地对波形求和
  • j发送节点的输出在被i节点根据函数attenuate(i,j)
  • 接收之前将被衰减

所以这是代码:

#create empty 2d array (no_rx_nodes x no_samples for each waveform)
waveforms = np.zeros((no_nodes, max(end_clock))) 

for i in range(no_nodes): #calculate the waveform for each receiver
    for j in range(no_nodes): #sum the waveforms produced by each transmitter
        waveforms[i, start_clock[j]:end_clock[j]] += output[j,:] * attenuate(i,j)
return waveforms

对上述内容的一些评论:

  • output[j, :]是发射器j的输出波形
  • waveforms[i,:]是接收方i收到的波形

我希望我在这里要完成的工作相当清楚。因为产生的波形非常大(大约10 ^ 6个样本),我也尝试将此代码转换为cython,但没有注意到任何特定的加速(可能是5-10倍更好但不多)。我想知道是否还有其他任何东西可以求助,以便加速,因为它是整个模拟的真正瓶颈(它需要计算几乎与其余代码一样多的时间,实际上是相当的比这更复杂。)

3 个答案:

答案 0 :(得分:5)

这是一个内存带宽限制问题,大约3GB / s内存带宽,你可以得到的最好的内部循环大约2-4ms。 达到该范围你需要阻止你的内部循环以更好地利用cpu缓存(numexpr为你做这个):

for i in range(no_nodes):
    for j in range(no_nodes):
        # should be chosen so all operands fit in the (next-to-)last level cache
        # first level is normally too small to be usable due to python overhead
        s  = 15000 
        a = attenuation[i,j]
        o = output[j]
        w = waveforms[i]
        for k in range(0, w.size, s): 
            u = min(k + s, w.size)
            w[k:u] += o[k:u] * a
        # or: numexpr.evaluate("w + o * a", out=w)

使用float32数据而不是float64也应该是内存带宽需求的一半,并且性能提高一倍。

要获得更大的加速,您必须重新设计完整的算法以获得更好的数据位置

答案 1 :(得分:4)

我认为内存访问会削弱您的性能,而且您可以做的很少。通过使用就地操作以及预分配和重用缓冲区,您可以稍微加快速度。作为代码的玩具示例,没有时间对齐:

def no_buffer(output, attenuate):
    waveforms = np.zeros_like(output)
    for i in xrange(len(output)):
        for j in xrange(len(output)):
            waveforms[i,:] += output[j, :] * attenuate[i, j]

    return waveforms

def with_buffer(output, attenuate):
    waveforms = np.zeros_like(output)
    buffer_arr = np.empty_like(output[0])
    for i in xrange(len(output)):
        for j in xrange(len(output)):
            np.multiply(output[j, :], attenuate[i, j], out=buffer_arr)
            np.add(waveforms[i, :], buffer_arr, out=waveforms[i, :])

    return waveforms

o = np.random.rand(20, 1e6)
a = np.random.rand(20, 20)

In [17]: np.allclose(no_buffer(o, a), with_buffer(o, a))
Out[17]: True

In [18]: %timeit no_buffer(o, a)
1 loops, best of 3: 2.3 s per loop

In [19]: %timeit with_buffer(o, a)
1 loops, best of 3: 1.57 s per loop

我猜这比没有好。

当然,如果你可以摆脱时间对齐的东西,你的操作只是一个矩阵乘法,最好让BLAS处理它。在我的系统上,使用MKL:

In [21]: np.allclose(with_buffer(o, a), np.dot(o.T, a.T).T)
Out[21]: True

In [22]: %timeit np.dot(o.T, a.T).T
10 loops, best of 3: 123 ms per loop

答案 2 :(得分:2)

仅仅为了实验,假设每个发射机的输出是时间对齐的,因此不需要时钟。我想出了一个使用大量广播的版本,因此完全消除了for循环。但是,它慢了3倍。 这是我写的代码:

import numpy as np
import time

def calc(no_nodes):

    output = np.random.rand(no_nodes, 7e5) #some random data, 7e5 samples here
    attenuate= np.random.rand(no_nodes,no_nodes) #some random data
    start_time = time.time()
    output_per_node = np.zeros((no_nodes,no_nodes,7e5))
    output_per_node += output[None, :, :]
    data = attenuate[:,:,None] * output_per_node
    waveforms = np.sum(data, axis=1)
    end_time = time.time()
    print end_time - start_time
    return waveforms

常规嵌套for循环实现将是:

def calc1(no_nodes):
    output = np.random.rand(no_nodes, 7e5)
    attenuation = np.random.rand(no_nodes,no_nodes)
    waveforms = np.zeros((no_nodes, 7e5))
    start_time = time.time()
    for i in range(no_nodes):
        for j in range(no_nodes):
            waveforms[i] += output[j] * attenuation[i,j]
    print time.time() - start_time
    return waveforms

那么这里到底发生了什么? Numpy应该是如此超快,你无法跟上。我不是说它通常不是,但在这个特殊的例子中,某些事情并不顺利。即使你将这两个代码转换为cython,第二个(带有for循环)比使用广播的代码要快得多。你觉得我在这里做错了什么? 注意:尝试使用no_nodes = 10

对于任何感兴趣的人,你可以找到ipython笔记本,上面的代码显示了性能上的差异,这里有ipynb和html格式:

非常感谢任何反馈。