numpy乘法矩阵保留第三轴 [英] numpy multiply matrices preserve third axis

问题描述

如果我有k许多由n,m,k numpy数组表示的n,m矩阵，如何将其乘以k许多由'm，j，k'numpy数组表示的k矩阵同时给我一个n,j,k ndarray?

if I have k many n,m matrices represented by a n,m,k numpy array, how can I multiply that by k many m,j matrices represented by a 'm,j,k' numpy array at the same time giving me a n,j,k ndarray?

换句话说，我需要执行k n,m * m,j = n,j的许多矩阵乘法.可以一次执行吗?

In other words, I need to perform k many matrix multiplications of n,m * m,j = n,j. Is it possible to perform them at once?

所有尺寸都不同，但通常都很大.

All of the dimensions vary, but are in general large.

推荐答案

@Ophion的第二种解决方案可以不使用循环，并且在使用较大尺寸时会更快:

The second solution of @Ophion can do without a loop, and it is faster with larger dimension:

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慢:

And slower than enisum when dimension is small:

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 s.

Sure, this is for python 2.x, in 3.x, be aware that the map returns map objects.

这篇关于numpy乘法矩阵保留第三轴的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！