使用numpy加快循環速度

Question

下一個for循環如何通過numpy獲得加速？ 我猜這里可以使用一些花哨的索引技巧，但我不知道哪一個（可以在這里使用einsum嗎？）。

a=0
for i in range(len(b)):
    a+=numpy.mean(C[d,e,f+b[i]])*g[i]

編輯： C是一個形狀類似於(20, 1600, 500) 20，1600，500）的numpy 3D數組。 d,e,f是“有趣”的點的索引（ d,e,f長度相同，大約900），b和g的長度相同（大約50）。 平均值取C所有點的索引為d,e,f+b[i]

Answer 1

您可以執行以下技巧：

C[d, e][:, np.add.outer(f, b)].dot(g).diagonal().mean()

通過過早采用將形成對角線的元素，可以進一步提高：

C[d, e][np.arange(len(d))[:, None], np.add.outer(f, b)].dot(g).mean()

Answer 2

它與循環版本非常相似：

np.sum(np.mean(C[d,e,f+b[:,None]], axis=1) * g)

您可以將求和和乘法組合成點積：

C[d,e,f+b[:,None]].mean(1).dot(g)

但是對於時間似乎並不重要； 索引操作是迄今為止最耗時的操作（至少在Numpy 1.8.0上）。 相比之下，原始代碼中的循環開銷微不足道。

Answer 3

時機

兩個會話都使用

In [1]: C = np.random.rand(20,1600,500)

In [2]: d = np.random.randint(0, 20, size=900)

In [3]: e = np.random.randint(1600, size=900)

In [4]: f = np.random.randint(400, size=900)

In [5]: b = np.random.randint(100, size=50)

In [6]: g = np.random.rand(50)

脾氣暴躁的1.9.0

In [7]: %timeit C[d,e,f + b[:,np.newaxis]].mean(axis=1).dot(g)
1000 loops, best of 3: 942 µs per loop

In [8]: %timeit C[d[:,np.newaxis],e[:, np.newaxis],f[:, np.newaxis] + b].mean(axis=0).dot(g)
1000 loops, best of 3: 762 µs per loop

In [9]: %%timeit                                               
   ...: a = 0
   ...: for i in range(len(b)):                                     
   ...:     a += np.mean(C[d, e, f + b[i]]) * g[i]
   ...: 
100 loops, best of 3: 2.25 ms per loop

In [10]: np.__version__
Out[10]: '1.9.0'

In [11]: %%timeit
(C.ravel()[np.ravel_multi_index((d[:,np.newaxis],
                                 e[:,np.newaxis],
                                 f[:,np.newaxis] + b), dims=C.shape)]
 .mean(axis=0).dot(g))
   ....: 
1000 loops, best of 3: 940 µs per loop

脾氣暴躁的1.8.2

In [7]: %timeit C[d,e,f + b[:,np.newaxis]].mean(axis=1).dot(g)
100 loops, best of 3: 2.81 ms per loop

In [8]: %timeit C[d[:,np.newaxis],e[:, np.newaxis],f[:, np.newaxis] + b].mean(axis=0).dot(g)
100 loops, best of 3: 2.7 ms per loop

In [9]: %%timeit                                               
   ...: a = 0
   ...: for i in range(len(b)):                                     
   ...:     a += np.mean(C[d, e, f + b[i]]) * g[i]
   ...: 
100 loops, best of 3: 4.12 ms per loop

In [10]: np.__version__
Out[10]: '1.8.2'

In [51]: %%timeit
(C.ravel()[np.ravel_multi_index((d[:,np.newaxis],
                                 e[:,np.newaxis],
                                 f[:,np.newaxis] + b), dims=C.shape)]
 .mean(axis=0).dot(g))
   ....: 
1000 loops, best of 3: 1.4 ms per loop

描述

您可以從一開始就使用坐標廣播技巧來充實您的50x900陣列：

In [158]: C[d,e,f + b[:, np.newaxis]].shape
Out[158]: (50, 900)

從這一點來看， mean和dot將使您到達目的地：

In [159]: C[d,e,f + b[:, np.newaxis]].mean(axis=1).dot(g)
Out[159]: 13.582349962518611

In [160]: 
a = 0
for i in range(len(b)):       
    a += np.mean(C[d, e, f + b[i]]) * g[i]
print(a)
   .....: 
13.5823499625

它比循環版本快3.3倍：

In [161]: %timeit C[d,e,f + b[:, np.newaxis]].mean(axis=1).dot(g)
1000 loops, best of 3: 585 µs per loop

In [162]: %%timeit                                               
a = 0
for i in range(len(b)):                                     
    a += np.mean(C[d, e, f + b[i]]) * g[i]
   .....: 
1000 loops, best of 3: 1.95 ms per loop

該陣列的大小很大，因此必須考慮CPU緩存。 我不能說我知道np.sum是如何遍歷數組的，但是在2d數組中，總是有一種更好的方法（當您選擇的下一個元素與內存相鄰時）和一種稍差的方法（當在數組中找到下一個元素時）大步前進）。 讓我們看看是否可以通過在索引期間轉置數組來贏得更多收益：

In [196]: C[d[:,np.newaxis], e[:,np.newaxis], f[:,np.newaxis] + b].mean(axis=0).dot(g)
Out[196]: 13.582349962518608

In [197]: %timeit C[d[:,np.newaxis], e[:,np.newaxis], f[:,np.newaxis] + b].mean(axis=0).dot(g)
1000 loops, best of 3: 461 µs per loop

比循環快4.2倍。

Answer 4

從結構上來講，您唯一希望的速度是使用以下代碼：

#Initialize a 4-D array
aggregated = numpy.zeros((len(d), len(e), len(f), len(b)))
#Populate it by the shifted copies of C
for i in range(len(b)):
    aggregated[:, :, :, i] = C[d, e, f + b[i]]

#Compute the mean on the first three axes
means = numpy.mean(aggregated, axis=(0, 1, 2))
#Multiply term-by-term by g (be careful that means and g have the same size!) and sum
a = numpy.sum(means * g)

但是，這不能保證計算會更快，甚至可能由於以下原因而變慢：

• 填充4-D陣列的成本不可忽略，因為它復制內存
• b非常小，因此您無論如何都不會贏。 如果b更大，只要d，e，f也變小，這就會變得很有趣。

無論如何，您都應該對兩種解決方案進行基准測試。 您也可以嘗試使用Cython之類的東西來執行for循環，但這似乎有點過頭了。