为什么我的矢量化Numpy代码比非矢量化代码耗时更长

Question

所以我用大量的数据计算泊松分布。我有一个形状数组（2666667,19）-“尖峰”，还有一个形状数组（19100）-“placefields”。我曾经有一个for循环，它遍历2666667维，大约需要60秒才能完成。然后，我了解到，如果我对循环进行矢量化，它会变得更快，所以我尝试这样做。矢量化表单工作并输出相同的结果，但现在需要120秒：

以下是原始循环（60秒）：

def compute_probability(spikes,placefields):
    nTimeBins = len(spikes[0])
    probability = np.empty((nTimeBins, 99)) #empty probability matrix
    for i in range(nTimeBins):
        nspikes = np.tile(spikes[:,i],(99))
        nspikes = np.swapaxes(nspikes,0,1)
        maxL = stats.poisson.pmf(nspikes,placefields)
        maxL = maxL.prod(axis=0)
        probability[i,:] = maxL
    return probability

这是矢量形式（120s）

def compute_probability(spikes,placefields):  

    placefields = np.reshape(placefields,(19,99,1))
    #prepared placefields

    nspikes = np.tile(spikes, (99,1,1))
    nspikes = np.swapaxes(nspikes,0,1)
    #prepared nspikes

    probability = stats.poisson.pmf(nspikes,placefields)
    probability = np.swapaxes(probability.prod(axis=0),0,1)
    return probability

为什么这么慢。我认为可能是矢量化形式创建的平铺数组太大了，它们占用了大量内存。我怎样才能让它跑得更快？ 下载samplespikes和sampleplacefields（如评论所建议的）-https://mega.nz/file/lpRF1IKI#YHq1HtkZ9EzYvaUdlrMtBwMg-0KEwmhFMYswxpaozXc

编辑： 问题是，尽管它是矢量化的，但庞大的阵列占用了太多的内存。我已经将计算分为几部分，现在它做得更好了：

placefields = np.reshape(placefields,(len(placefields),99,1))
nspikes = np.swapaxes(np.tile(spikes, (xybins,1,1)),0,1)
probability = np.empty((len(spikes[0]), xybins))

chunks = len(spikes[0])//20
n = int(len(spikes[0])/chunks)
for i in range(0,len(nspikes[0][0]),n):
    nspikes_chunk = nspikes[:,:,i:i+n]
    probability_chunk = stats.poisson.pmf(nspikes_chunk,placefields)
    probability_chunk = np.swapaxes(probability_chunk.prod(axis=0),0,1)
    if len(probability_chunk)<(len(spikes)//chunks):
        probability[i:] = probability_chunk
    else:
        probability[i:i+len(probability_chunk)] = probability_chunk

Answer 1

这可能是由于内存/缓存效应造成的

第一个代码处理适合CPU缓存的小型阵列。这并不好，因为每个Numpy函数调用都需要一些时间。第二个代码修复了这个问题。但是，它在多个GiB的内存中分配/填充巨大的数组。在CPU缓存中工作要比在主存（RAM）中快得多。当工作数组只使用一次时（因为昂贵的操作系统页面错误），情况尤其如此，这在代码中似乎就是如此。如果内存不足，操作系统将在SSD/HDD存储设备中读取/写入临时数据，与RAM和CPU缓存相比，这些设备的速度非常慢

最好的解决方案可能是处理块，这样操作既可以矢量化（减少Numpy函数调用的开销），又可以放入CPU缓存（减少RAM读/写的成本）。请注意，在主流PC处理器上，最后一级缓存的大小通常只有很少的MiB

一个重要的信息是，矢量化并不总是让事情变得更快。为了获得更好的性能，应该关注被操纵数据块的大小，以便它们适合CPU缓存

PS：请注意，如果您不太关心精度，可以使用简单精度（np.float32）而不是双精度（np.float64）来加快计算速度