加速numpy嵌套循环

Question

加速numpy嵌套循环

4

我正在使用numpy和cython在python中编写一个无线网络模拟程序，假设有一些节点no_nodes在二维平面上随机分布，它们发送一些波形和相应的接收器，同样随机分布在二维平面上。每个发送节点产生一个我称之为output的波形（每个节点可以产生不同长度的输出）。

我的目标是将每个节点的输出加总成一个大波形，作为每个接收器进行解调等操作的输入。现在有两个关键点：

发射器异步发送，因此必须对每个发射器维护一个start_clock和一个end_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倍，但不会更多）。我想知道是否还有其他可以使用的方法来获得速度提升，因为它是整个模拟的真正瓶颈（计算时间几乎与其余代码一样长，实际上比其余代码更复杂）。

- user113478

代码看起来相当高效，除非“no_nodes”很大。你确定瓶颈不在“attenuate”函数中吗？ - loopbackbee

no_nodes 最多只有10-20个。attenuate 只是从距离矩阵中简单地进行乘法运算。 - user113478

1

@user113478，所以你只需要计算矩阵 attenuate = (attn_const/distances) ** attn_expon，这将产生与 distances 相同大小的数组，对吗？ - Henry Gomersall

2

哦，在原来的问题中也发帖。我的意思是，展示尽可能多的代码。你的操作是线性的，因此如果从算法的角度讲，总是可以重新排序。你的代码是在一个点计算所有绿函数的总和吗？如果是这样的话，你从哪里得到了这个算法？在线性情况下（看起来你已经有了），我理解这应该在 Fourier 领域中解决最好。 - Henry Gomersall

2

@user113478，嗯，通常在傅里叶域中执行卷积是相当典型的。问题是，你正在寻求关于算法某个小方面的帮助，而我认为最好的策略是重新设计整个算法。我的建议是离开并认真思考算法。 - Henry Gomersall

显示剩余16条评论

3个回答

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

- Jaime

如果问题确实与数组大小有关，而且您不能使用“dot”，那么研究一下“numexpr”可能是值得的。 - Midnighter

1

@Midnighter 你最好写几个简单的循环，逐块完成点积操作。 - Daniel

@Jaime，我以为我懂线性代数，但是在看了你的回答之后，我改变了我的想法！非常好的答案，非常感谢！ - user113478

@Jaime，如果我需要将output[j,:]与attenuate[i,j,:]进行卷积，其中attenuate将是一个3D数组（no_nodes x no_nodes x no_samples_for_channel_model）（假设attenuate是表示信号从发射机j到达接收器i所经过的信道模型）。你能提出解决方法吗？是否可以使用np.dot进行操作？ - user113478

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格式均可：

HTML格式的IPython笔记本（无法编辑）
.ipynb格式的IPython笔记本（供您使用）

非常感谢您的任何反馈。

- user113478

我喜欢你的方法，你是否使用np.allclose()检查了相等性呢？例如，看看一个方法是否与另一个方法给出相同的结果。你应该将这个最后的问题作为一个新答案发布，这将引发良好的讨论... - Saullo G. P. Castro

@SaulloCastro，是的，它们返回完全相同，实际上我将在几分钟内发布一个IPython笔记本，供任何感兴趣的人使用。 - user113478

1

在您使用广播的解决方案中，您正在创建一个比其他解决方案中使用的任何数组大20倍的数组。我认为这就是减速的原因所在。如果您真的可以不需要时间对齐，请仔细检查您的代码：您已经重新实现了矩阵乘法，而np.dot可以比您快很多倍... - Jaime

@Jaime，你能详细说明一下吗？如何使用np.dot来完成这个操作？ - user113478

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- jtaylor · Accepted Answer

这是一个内存带宽受限的问题，大约有3GB/s的内存带宽，你能得到的最佳结果是内部循环大约2-4毫秒。为了达到这个限制，你需要阻塞你的内部循环，更好地利用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，可以将内存带宽要求减半并将性能提高一倍。

要获得更大的加速效果，需要重新设计整个算法以获得更好的数据本地性。