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

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μs。

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 2

使用正确重构的数组上的爱因斯坦求和，这是相对简单的：

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)

（注意：除了帮助理解爱因斯坦总结的工作原理之外，没有理由像我之前那样重塑扁平阵列（在我看来，比（3）更容易看到（1x3）矩阵）矩阵。）所以，我已经把它删除了。）

爱因斯坦惯例意味着对每个操作数重复的索引求和。在我们的两个矩阵具有a.shape = (n,)和b.shape = (n,...)形状的一般情况下，我们希望仅对a和b的第一维进行求和。我们不关心b中其他维度的深度，或者可能有多少维度的深度，因此我们使用...作为剩余维度的全部内容。求和维从输出数组中消失，因此输出数组包含所有其他维（即...）。

传递给einsum的下标字符串会捕获所有这些信息。在字符串的输入端（->左侧的所有内容），我们标记每个操作数的索引（即输入矩阵a和b），用逗号分隔。重复指出的指数（即i）。在字符串的输出端（->的右侧），我们指示输出索引。我们的函数不需要输出字符串，因为我们想输出总和中未包含的所有维度（我认为）。

Answer 3

对于多维数组的这种总和减少，我认为我们可以在重塑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

反正有没有让这个numpy阵列操作更快？

3 个答案: