Python:在循环中矢量化矩阵乘法?
我有一个N-by-M数组,在每个数组的入口,我需要做一些NumPy操作并将结果放在那里 现在,我正以一种天真的方式通过一个双循环来实现:Python:在循环中矢量化矩阵乘法?,python,numpy,vectorization,Python,Numpy,Vectorization,我有一个N-by-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,
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操作,我认为一定有一种方法可以更快或无循环地完成。我们可以使用
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
来利用多核
-
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})
另一个非常简单的选择是:
result = (values1 * values2 * weight * is_relevant[:, :, np.newaxis, np.newaxis]).sum((2, 3, 4))
最后一个解决方案比这个更快。比较时间:
%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)
..
(np.all(out==result)-->False
@wwii我们处理的是浮动pt数。因此,最好使用np.allclose()
。每次都会收到我的消息。@sibbsbsbgambing请查看结尾处的编辑-非常大的数组
。想知道你可能会得到什么样的加速效果。MKL
也应该会有帮助。如果你有任何提升,请告诉我。@sibbsbsbgambing这可能相关-?如果没有答案,请随时询问一个新的.