Question

如果我有k个n,m numpy数组所代表的n,m,k个矩阵，我怎样才能将其乘以k多个m,j矩阵所代表的矩阵＆＃39;米，J，K＆＃39; numpy数组同时给我一个n,j,k ndarray？

换句话说，我需要执行k n,m * m,j = n,j的许多矩阵乘法。是否可以立即执行它们？

编辑：所有尺寸都有所不同，但一般都很大。

Answer 1

这实际上取决于阵列的大小和形状。如果n，m和j很小，您可以执行以下操作：

import numpy as np
>>> a = np.random.rand(5,2,1E6)
>>> b = np.random.rand(2,5,1E6)
>>> out = np.einsum('nmk,mjk->njk',a,b)
>>> out.shape
(5, 5, 1000000)

如果n，m和j很大，您可能希望利用这样的BLAS：

>>> a= np.random.rand(1E3,1E2,5)
>>> b= np.random.rand(1E2,1E3,5)
>>> out = np.empty((1E3,1E3,5))
>>> for x in range(out.shape[-1]):
...     out[:,:,x] = np.dot(a[:,:,x], b[:,:,x])

请记住，numpy数组是行主要的。您可能希望根据您的操作重新排列数据。

Answer 2

@Ophion的第二个解决方案可以在没有循环的情况下进行，并且在更大的维度下更快：

In [65]:

#k,n,m,j=2,4,5,6
k,n,m,j=100,101,102,103
A=np.random.random((n,m,k))
B=np.random.random((m,j,k))
In [66]:

%timeit np.rollaxis(np.array(map(np.dot, np.rollaxis(A,2), np.rollaxis(B,2))), 0, 3)
1 loops, best of 3: 313 ms per loop
In [67]:

%timeit np.einsum('nmk,mjk->njk',A,B)
1 loops, best of 3: 793 ms per loop

尺寸较小时比enisum慢：

In [68]:

k,n,m,j=2,4,5,6
#k,n,m,j=100,101,102,103
A=np.random.random((n,m,k))
B=np.random.random((m,j,k))
In [69]:

%timeit np.rollaxis(np.array(map(np.dot, np.rollaxis(A,2), np.rollaxis(B,2))), 0, 3)
10000 loops, best of 3: 73.7 µs per loop
In [70]:

%timeit np.einsum('nmk,mjk->njk',A,B)
100000 loops, best of 3: 13.5 µs per loop

当然，这是针对python 2.x的，3.x，请注意map会返回map object。

numpy multiply矩阵保留第三轴

2 个答案: