numpy 在python中连接数组有更快的方法吗?

ekqde3dh  于 2023-08-05  发布在  Python
关注(0)|答案(1)|浏览(106)

我有一个函数,它执行以下操作:

  • 获取一个数组,然后对该数组的切片执行任意操作,生成新数组
  • 将所有新数组连接成一个数组并返回它。

一个非常直接的方法是使用np.concatenate编写以下代码:

import numpy as np
from numba import njit

arr = np.ones(10000)

@njit(fastmath=True)
def EQ1(a):
    
    eq1 = 2*a[1:-1]
    eq2 = 4*a[0:2]
    
    sys = (eq1,eq2)
    
    sys = np.concatenate(sys)
    
    return sys

字符串
然后我决定寻找更快的答案,奇怪的是,执行一个numba jitted函数将每个单独的元素附加到预先分配的数组中要快得多:

@njit(fastmath=True)
def EQ2(a):
    
    sys = np.empty_like(a)
    l = 0
    for i in range(1,len(a)-1):
        eq1 = 2*a[i]
        sys[l] = eq1
        l+=1
   
    for i in range(2):
        eq2 = 4*a[i]
        sys[l] = eq2
        l+=1
    
    return sys

%timeit EQ1(arr) 6.31 µs ± 9.72 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit EQ2(arr) 3.48 µs ± 2.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)


但是第二种方法随着eq数量的增加开始占用大量空间。所以我想知道有没有更快的替代品。我已经尝试了append(),extend(),但for循环jitted仍然更快。谢啦,谢啦

kg7wmglp

kg7wmglp1#

你不是在问连锁速度。相反,为了完成一项任务,你会问“我需要多少次内存操作?”“目标是懒惰,尽可能少做。
我们在这里基本上不做任何计算这完全是一个专注于移动记忆块的练习。

公式1

这些操作的成本可以忽略不计,不再考虑:

eq2 = 4 * a[0:2]
    sys = (eq1, eq2)

字符串
我们将把“写一堆1”arr = np.ones(10_000)作为“一个内存操作”,一个数组的写入量通过总线到达DRAM。我相信这是以一种合理的方式发生的,每个缓存行只被写入一次。这就是我们的计时单位。

eq1 = 2 * a[1:-1]


在这里,我们读取一个单元,由于缓存,它在运行多次的基准测试上下文中似乎具有零成本。我们将每个项目平凡地加倍,并以一个单位的成本将其发送到临时变量。

sys = np.concatenate(sys)


现在,concatenate以零成本读取一个(缓存)单元,并产生额外的写入单元到新分配的存储的开销。
总计:两个单位。

EQ2

range(2)循环的成本可以忽略不计,因此不再进一步考虑。类似于.empty_like()

for i in range(1, len(a) - 1):
        eq1 = 2 * a[i]
        sys[l] = eq1
        l += 1


从缓存中阅读a是免费的。存储两倍于sys的数据会花费一个单位。相对于内存操作,乘以2不是这个流水线中的瓶颈。在寄存器中缓冲一段时间,而不是在DRAM中缓冲一段时间,是一个巨大的胜利。
因此,EQ1的临时变量花费了两个单位的内存操作,而EQ2只花费了一个单位。
这个故事的寓意:“大”临时变量是昂贵的,因为它们不是缓存友好的。

相关问题