有没有办法让这个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）。在字符串的输出端（在->的右边），我们指示输出索引。我们的函数不需要输出字符串，因为我们想要输出所有未包含在求和中的维度（我认为）。