Python：循环中的向量化矩阵乘法？答案

【问题标题】：Python: Vectorizing Matrix Multiplications in the Loops?Python：循环中的向量化矩阵乘法？
【发布时间】：2019-05-18 02:21:10
【问题描述】：

我有一个 N×M 数组，我需要在其中的每个条目处执行一些 NumPy 操作并将结果放在那里。

现在，我用双循环以天真的方式来做：

import numpy as np

N = 10
M = 11
K = 100

result = np.zeros((N, M))

is_relevant = np.random.rand(N, M, K) > 0.5
weight = np.random.rand(3, 3, K)
values1 = np.random.rand(3, 3, K)
values2 = np.random.rand(3, 3, K)

for i in range(N):
    for j in range(M):
        selector = is_relevant[i, j, :]
        result[i, j] = np.sum(
            np.multiply(
                np.multiply(
                    values1[..., selector],
                    values2[..., selector]
                ), weight[..., selector]
            )
        )

由于所有的循环内操作都是简单的 NumPy 操作，我认为必须有一种方法可以更快或无循环地做到这一点。

【问题讨论】：

标签： python numpy vectorization

【解决方案1】：

我们可以使用np.einsum 和np.tensordot 的组合-

a = np.einsum('ijk,ijk,ijk->k',values1,values2,weight)
out = np.tensordot(a,is_relevant,axes=(0,2))

或者，使用一个einsum 呼叫 -

np.einsum('ijk,ijk,ijk,lmk->lm',values1,values2,weight,is_relevant)

还有 np.dot 和 einsum -

is_relevant.dot(np.einsum('ijk,ijk,ijk->k',values1,values2,weight))

另外，通过将np.einsum 中的optimize 标志设置为True 来使用BLAS。

时间安排 -

In [146]: %%timeit
     ...: a = np.einsum('ijk,ijk,ijk->k',values1,values2,weight)
     ...: out = np.tensordot(a,is_relevant,axes=(0,2))
10000 loops, best of 3: 121 µs per loop

In [147]: %timeit np.einsum('ijk,ijk,ijk,lmk->lm',values1,values2,weight,is_relevant)
1000 loops, best of 3: 851 µs per loop

In [148]: %timeit np.einsum('ijk,ijk,ijk,lmk->lm',values1,values2,weight,is_relevant,optimize=True)
1000 loops, best of 3: 347 µs per loop

In [156]: %timeit is_relevant.dot(np.einsum('ijk,ijk,ijk->k',values1,values2,weight))
10000 loops, best of 3: 58.6 µs per loop

非常大的数组

对于非常大的数组，我们可以利用numexpr 来利用multi-cores -

import numexpr as ne

a = np.einsum('ijk,ijk,ijk->k',values1,values2,weight)
out = np.empty((N, M))
for i in range(N):
    for j in range(M):
        out[i,j] = ne.evaluate('sum(is_relevant_ij*a)',{'is_relevant_ij':is_relevant[i,j], 'a':a})

【讨论】：

..(np.all(out == result) --> False
@wwii 我们正在处理浮动点数。所以，最好使用np.allclose()。
每次都能得到我。
@SibbsGambling 请检查最后的编辑 - Very large arrays。想知道你可能会得到什么样的加速。 MKL 也应该有所帮助。让我知道你是否有任何提升。
@SibbsGambling 这可能是相关的 - stackoverflow.com/questions/50295180?如果没有回答，请随意问一个新的。

【解决方案2】：

另一个非常简单的选择就是：

result = (values1 * values2 * weight * is_relevant[:, :, np.newaxis, np.newaxis]).sum((2, 3, 4))

Divakar's 最后一个解决方案比这更快。比较时间：

%timeit np.tensordot(np.einsum('ijk,ijk,ijk->k',values1,values2,weight),is_relevant,axes=(0,2))
# 30.9 µs ± 1.71 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit np.einsum('ijk,ijk,ijk,lmk->lm',values1,values2,weight,is_relevant)
# 379 µs ± 486 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit np.einsum('ijk,ijk,ijk,lmk->lm',values1,values2,weight,is_relevant,optimize=True)
# 145 µs ± 1.89 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit is_relevant.dot(np.einsum('ijk,ijk,ijk->k',values1,values2,weight))
# 15 µs ± 124 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit (values1 * values2 * weight * is_relevant[:, :, np.newaxis, np.newaxis]).sum((2, 3, 4))
# 152 µs ± 1.4 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

【讨论】：