Question

我有一个代码在numpy数组上运行操作。虽然线性代数运算似乎很快，但我现在发现了一个不同问题的瓶颈：两个不同数组的总和。在下面的示例中，WE3和T1是两个1000X1000X1000数组。首先，我使用numpy操作计算WE3，然后对这些数组求和。

import numpy as np
import scipy as sp
import time
N = 100
n = 1000
X = np.random.uniform(size = (N,n))
wE = np.mean(X,0)
wE3 = np.einsum('i,j,k->ijk', wE, wE, wE) #22 secs
T1 = np.random.uniform(size = (n,n,n))
a = wE3 + T1 #115 secs

wE3的计算需要22秒，而WE3和T1之间的加法需要115秒。

为什么这些数组的总和比WE3的计算慢得多？它们应该或多或少具有相同的复杂性。

有没有办法加快代码的速度？

Answer 1

np.einsum('i,j,k->ijk', wE, wE, wE)部分没有进行任何总和减少，基本上只是广播元素乘法。所以，我们可以用这样的东西代替它 -

wE[:,None,None] * wE[:,None] * wE

运行时测试 -

In [9]: # Setup inputs at 1/5th of original dataset sizes
   ...: N = 20
   ...: n = 200
   ...: X = np.random.uniform(size = (N,n))
   ...: wE = np.mean(X,0)
   ...: 

In [10]: %timeit np.einsum('i,j,k->ijk', wE, wE, wE)
10 loops, best of 3: 45.7 ms per loop

In [11]: %timeit wE[:,None,None] * wE[:,None] * wE
10 loops, best of 3: 26.1 ms per loop

接下来，我们有wE3 + T1，其中T1 = np.random.uniform(size = (n,n,n))看起来并没有大的帮助，因为无论如何我们必须创建T1然后＃39} ; s只是元素添加。我们似乎可以使用np.add让我们将结果写回其中一个数组：wE3或T1。我们假设我们选择T1，如果可以修改的话。我想这会带来轻微的内存效率，因为我们不会在工作空间中添加另一个变量。

因此，我们可以做 -

np.add(wE3,T1,out=T1)

运行时测试 -

In [58]: def func1(wE3):
    ...:     T1 = np.random.uniform(size = (n,n,n))
    ...:     return wE3 + T1
    ...: 
    ...: def func2(wE3):
    ...:     T1 = np.random.uniform(size = (n,n,n))
    ...:     np.add(wE3,T1,out=T1)
    ...:     return T1
    ...: 

In [59]: # Setup inputs at 1/4th of original dataset sizes
    ...: N = 25
    ...: n = 250
    ...: X = np.random.uniform(size = (N,n))
    ...: wE = np.mean(X,0)
    ...: wE3 = np.einsum('i,j,k->ijk', wE, wE, wE)
    ...: 

In [60]: %timeit func1(wE3)
1 loops, best of 3: 390 ms per loop

In [61]: %timeit func2(wE3)
1 loops, best of 3: 363 ms per loop

使用@Aaron's suggestion，我们可以使用循环并假设将结果写回wE3是可以的，我们可以这样做 -

wE3 = wE[:,None,None] * wE[:,None] * wE
for x in wE3: 
    np.add(x, np.random.uniform(size = (n,n)), out=x)

最终结果

因此，回顾所有建议的改进，最后运行时测试结果是 -

In [97]: def func1(wE):
    ...:     wE3 = np.einsum('i,j,k->ijk', wE, wE, wE)
    ...:     T1 = np.random.uniform(size = (n,n,n))
    ...:     return wE3 + T1
    ...: 
    ...: def func2(wE):
    ...:     wE3 = wE[:,None,None] * wE[:,None] * wE
    ...:     for x in wE3: 
    ...:         np.add(x, np.random.uniform(size = (n,n)), out=x)
    ...:     return wE3
    ...: 

In [98]: # Setup inputs at 1/3rd of original dataset sizes
    ...: N = 33
    ...: n = 330
    ...: X = np.random.uniform(size = (N,n))
    ...: wE = np.mean(X,0)
    ...: 

In [99]: %timeit func1(wE)
1 loops, best of 3: 1.09 s per loop

In [100]: %timeit func2(wE)
1 loops, best of 3: 879 ms per loop

Answer 2

是否有任何已知的原因导致这些数组的总和比WE3的计算慢？

数组wE3，T1和a每个都需要8千兆字节的内存。您的物理内存可能已用完，swap memory访问会导致您的性能下降。

有没有办法加快代码的速度？

获得更多物理内存（即RAM）。

如果无法做到这一点，请查看您要对这些阵列执行的操作，并查看是否可以批量工作，以便处理批处理时所需的总内存保持在物理内存的限制范围内。

Answer 3

你应该真的使用Numba的Jit（只是及时编译器）。它是一条纯粹的numpy管道，非常适合Numba。

您只需将上面的代码放入一个函数中，然后将@jit decorater放在顶部。它获得了接近Cython的加速。

但是，正如其他人所指出的那样，您似乎正在尝试使用对于本地计算机来说太大的数据，并且numba无法解决您的问题

加快大阵列之间的numpy总结

3 个答案: