为什么我的Numpy向量化代码运行速度比非向量化循环代码更慢？

我在使用大量数据计算泊松分布，现有形状为(2666667,19)的数组spikes，以及形状为(19,100)的数组placefields。最初编写的遍历2666667维度的for循环代码运行耗时约60秒，后续我对代码做了向量化改造，改造后的代码可正常输出正确结果，但耗时反而升至120秒。

原循环版本代码（耗时60s） #

def compute_probability(spikes,placefields):
    nTimeBins = len(spikes[0])
    probability = np.empty((nTimeBins, 99)) # 初始化概率矩阵
    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))
    # 预处理placefields

    nspikes = np.tile(spikes, (99,1,1))
    nspikes = np.swapaxes(nspikes,0,1)
    # 预处理nspikes

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

问题定位与优化方案 #

向量化后性能下降的原因已找到：虽然代码完成了向量化，但生成的超大数组占用了过多内存，导致运行效率降低。将计算拆分为分块执行后性能明显提升，优化后代码如下：

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

内容的提问来源于stack exchange，提问作者Alex Markov