有沒有辦法讓這個numpy數組操作更快？

Question

假設我有一個矩陣：

A = [[2, 1] 
    [1, 2]] 

和矩陣列表：

B = [[1, 0] C = [[2, 1], D = [[0, 0], E = [[1, 0], 
    [1, 0]]  [0, 0]]  [0, 0]]  [0, 0]] 

我首先要拉平A.flatten() = [2 1 1 2]然後把這些元素乘以B，C，D和E總和分別。所以：

A[0] * B + A[1]*C + A[2]*D + A[3]*E 

現在考慮一個更一般的情況：

A[0] * X_1 + A[1] * X_2 + ... + A[n-1] * X_n 

哪裏X_n可以有任何尺寸。這是我想出來的代碼：

import numpy as np 
from functools import reduce 
from operator import mul 

def product(iterable): 
    return reduce(mul, iterable) 


def create_table(old_shape, new_shape): 
    # Create X_1, X_2, ..., X_n 
    lookup = [] 
    for _ in range(product(old_shape)): 
     lookup.append(np.random.rand(*new_shape)) 
    return lookup 


def sum_expansion(arr, lookup, shape): 
    # A[0] * X_1 + ... + A[n-1] * X_n 
    new_arr = np.zeros(shape) 
    for i, a in enumerate(arr.flatten()): 
     new_arr += a * lookup[i] 

    return new_arr 

if __name__ == '__main__': 
    lookup = create_table((2, 2), (3, 3, 3)) 
    # Generate random 2 x 2 matrices. 
    randos = (np.random.rand(2, 2) for _ in range(100000)) 
    results = map(lambda x: sum_expansion(x, lookup, (3, 3, 3)), randos) 
    print(list(results)) 

要執行此代碼需要大約74秒在我的機器上。有什麼辦法可以減少這段代碼的時間？

Answer 1

對於這樣的總和，減少對多維數組，我認爲我們可以重塑randos最後兩軸合併成一個，像這樣後提示np.tensordot -

np.tensordot(np.array(randos).reshape(-1,4),lookup, axes=((-1),(0))) 

這裏還有一個與重塑第二陣列，而不是用於再次使用具有np.tensordot -

lookup_arr = np.asarray(lookup).reshape(2,2,3,3,3) 
out = np.tensordot(randos,lookup_arr,axes=((-2,-1),(0,1))) 

---

運行試驗 -

In [69]: randos = [np.random.rand(2, 2) for _ in range(100)] 

In [73]: lookup = create_table((2, 2), (3, 3, 3)) 

In [74]: lookup_arr = np.asarray(lookup).reshape(2,2,3,3,3) 

In [75]: out1 = np.tensordot(np.array(randos).reshape(-1,4),lookup, axes=((-1),(0))) 
    ...: out2 = np.tensordot(randos,lookup_arr,axes=((-2,-1),(0,1))) 
    ...: 

In [76]: np.allclose(out1, out2) 
Out[76]: True 

In [77]: %timeit np.tensordot(np.array(randos).reshape(-1,4),\ 
             lookup, axes=((-1),(0))) 
10000 loops, best of 3: 37 µs per loop 

In [78]: %timeit np.tensordot(randos,lookup_arr,axes=((-2,-1),(0,1))) 
10000 loops, best of 3: 33.3 µs per loop 

In [79]: %timeit np.asarray(lookup).reshape(2,2,3,3,3) 
100000 loops, best of 3: 2.18 µs per loop 

Answer 2

In [20]: randos = [np.random.rand(2, 2) for _ in range(10)] 

In [21]: timeit [sum_expansion(x,lookup,(3,3,3)) for x in randos]              10000 loops, best of 3: 184 µs per loop 

那把時間看起來不錯。每次調用sum_expansion需要18微秒。

In [22]: timeit create_table((2,2),(3,3,3))                    
100000 loops, best of 3: 14.1 µs per loop  

需要更多的時間來理解你在做什麼。我看到了很多Python迭代，以及很少的編碼numpy。

---

我得到使用einsum做乘法和總和提高了3倍：

def ein_expansion(arr, lookup, shape):                      
    return np.einsum('ij,ij...',arr, lookup) 

In [45]: L = np.array(lookup).reshape(2,2,3,3,3) 

In [43]: timeit [ein_expansion(r, L,(3,3,3)) for r in randos]               
10000 loops, best of 3: 58.3 µs per loop 

我們可以通過在多個randos陣列的一次操作得到進一步改進。

In [59]: timeit np.einsum('oij,ij...->o...',np.array(randos),L)               
100000 loops, best of 3: 15.8 µs per loop 

In [60]: np.einsum('oij,ij...->o...',np.array(randos),L).shape               
Out[60]: (10, 3, 3, 3) 

Answer 3

這是相對簡單的使用上正確整形陣列的愛因斯坦求和：

import numpy as np 


def do_sum(x, mat_lst): 
    a = np.array(x).flatten().reshape(1, -1) 
    print('A shape: ', a.shape) 
    b = np.stack(mat_lst) 
    print('B shape: ', b.shape) 
    return np.einsum('ij,jkl->kl', a, b) 

A = [[1,2],[3,4]] 
B = [[[1,1],[1,1]],[[2,2],[2,2]],[[3,3],[3,3]],[[4,4],[4,4]]] 

do_sum(A,B) 

輸出

A shape: (1, 4) 
B shape: (4, 2, 2) 

[[30 30] 
[30 30]] 

---

編輯 - 廣義情況

這是一個n-d輸入數組的列表。唯一的先決條件是x中元素的數量應等於mat_lst的長度。

def do_sum(x, mat_lst): 
    a = np.array(x).flatten() 
    b = np.stack(mat_lst) 
    print("A shape: {}\nB shape: {}".format(a.shape, b.shape)) 
    return np.einsum('i,i...', a, b) 

A = [[1,2],[3,4]] 
B = [np.random.rand(2,2,2) for _ in range(4)] 
do_sum(A,B) 

（注：沒有理由重塑扁平陣列，像我一樣以前，除了瞭解求和如何愛因斯坦的工作（在我看來，以幫助，它更容易可視化（1×3）矩陣比（3，）矩陣）所以，我已經在這裏刪除了。）

愛因斯坦約定意味着對每個操作數重複的索引進行求和。在我們的形狀爲a.shape = (n,)和b.shape = (n,...)的兩個矩陣的一般情況下，我們只想總結第一維a和b。我們不關心b中其他維度的深度，或者可能有多少維度，因此我們使用...作爲剩餘維度的全部內容。總和維度從輸出數組中消失，所以輸出數組包含所有其他維度（即...）。

傳遞到einsum的下標字符串捕獲所有這些信息。在字符串的輸入端（->左邊的所有內容），我們標記每個操作數的索引（即輸入矩陣a和b），用逗號分隔。重複索引（即i）。在字符串的輸出端（在->的右邊），我們指示輸出索引。我們的函數不需要輸出字符串，因爲我們想要輸出所有未包含在求和中的維度（我認爲）。