沿指定轴的两个3D矩阵之间的np.dot积 [英] np.dot product between two 3D matrices along specified axis
问题描述
我有两个3D矩阵:
a = np.random.normal(size=[3,2,5])
b = np.random.normal(size=[5,2,3])
我想要每个切片分别沿2和0轴的点积:
I want the dot product of each slice along 2 and 0 axes respectively:
c = np.zeros([3,3,5]) # c.size is 45
c[:,:,0] = a[:,:,0].dot(b[0,:,:])
c[:,:,1] = a[:,:,1].dot(b[1,:,:])
...
我想使用np.tensordot做到这一点(为了提高效率和速度) 我已经尝试过:
I would like to do that using np.tensordot (for efficiency and speed) I have tried:
c = np.tensordot(a, b, axes=[2,0])
但是我得到了一个包含36个元素(而不是45个)的4D数组. c.shape,c.size =((3L,2L,2L,3L),36).我在这里发现了类似的问题( Numpy张量:张量的张量张量),但这并不是我想要的,而且我无法推断该解决方案来解决我的问题. 总而言之,我可以使用np.tensordot计算上面显示的c数组吗?
but I get a 4D array with 36 elements (instead of 45). c.shape, c.size = ((3L, 2L, 2L, 3L), 36). I have found a similar question here (Numpy tensor: Tensordot over frontal slices of tensor) but it's not exactly what I want, and I was unable to extrapolate that solution to my problem. To summarise, can I use np.tensordot to compute c array show above?
更新#1
@hpaulj的答案是我想要的,但是在我的系统(python 2.7和np 1.13.3)中,这些方法相当缓慢:
The answer by @hpaulj is what I wanted, however in my system (python 2.7 and np 1.13.3) those aproaches are pretty slow:
n = 3000
a = np.random.normal(size=[n, 20, 5])
b = np.random.normal(size=[5, 20, n])
t = time.clock()
c_slice = a[:,:,0].dot(b[0,:,:])
print('one slice_x_5: {:.3f} seconds'.format( (time.clock()-t)*5 ))
t = time.clock()
c = np.zeros([n, n, 5])
for i in range(5):
c[:,:,i] = a[:,:,i].dot(b[i,:,:])
print('for loop: {:.3f} seconds'.format(time.clock()-t))
t = time.clock()
d = np.einsum('abi,ibd->adi', a, b)
print('einsum: {:.3f} seconds'.format(time.clock()-t))
t = time.clock()
e = np.tensordot(a,b,[1,1])
e1 = e.transpose(0,3,1,2)[:,:,np.arange(5),np.arange(5)]
print('tensordot: {:.3f} seconds'.format(time.clock()-t))
a = a.transpose(2,0,1)
t = time.clock()
f = np.matmul(a,b)
print('matmul: {:.3f} seconds'.format(time.clock()-t))
推荐答案
与tensordot
相比,使用einsum
更容易.因此,让我们从这里开始:
It's easier to work with einsum
than tensordot
. So let's start there:
In [469]: a = np.random.normal(size=[3,2,5])
...: b = np.random.normal(size=[5,2,3])
...:
In [470]: c = np.zeros([3,3,5]) # c.size is 45
In [471]: for i in range(5):
...: c[:,:,i] = a[:,:,i].dot(b[i,:,:])
...:
In [472]: d = np.einsum('abi,ibd->iad', a, b)
In [473]: d.shape
Out[473]: (5, 3, 3)
In [474]: d = np.einsum('abi,ibd->adi', a, b)
In [475]: d.shape
Out[475]: (3, 3, 5)
In [476]: np.allclose(c,d)
Out[476]: True
我不得不考虑一下尺寸.它有助于将a[:,:,i]
集中为2d,而对于b[i,:,:]
同样.因此,dot
总和位于两个数组的中间维度(大小2)上.
I had to think a bit about to match up the dimensions. It helped to focus on a[:,:,i]
as 2d, and similarly for b[i,:,:]
. So the dot
sum is over the middle dimension of both arrays (size 2).
在测试思路中,如果c
的前两个维度不同,可能会有所帮助.混合它们的机会会更少.
In testing ideas it might help if the first 2 dimensions of c
were different. There'd be less chance of mixing them up.
在tensordot
中指定dot
求和轴(轴)很容易,但是更难于约束其他尺寸的处理.这就是为什么要获得4D阵列的原因.
It's easy to specify the dot
summation axis (axes) in tensordot
, but harder to constrain the handling of the other dimensions. That's why you get a 4d array.
我可以使它与转置一起工作,然后取对角线:
I can get it to work with a transpose, followed by taking the diagonal:
In [477]: e = np.tensordot(a,b,[1,1])
In [478]: e.shape
Out[478]: (3, 5, 5, 3)
In [479]: e1 = e.transpose(0,3,1,2)[:,:,np.arange(5),np.arange(5)]
In [480]: e1.shape
Out[480]: (3, 3, 5)
In [481]: np.allclose(c,e1)
Out[481]: True
我计算出的值比需要的多得多,并且将其中的大部分扔掉了.
I've calculated a lot more values than needed, and thrown most of them away.
matmul
进行一些移调可能会更好.
matmul
with some transposing might work better.
In [482]: f = a.transpose(2,0,1)@b
In [483]: f.shape
Out[483]: (5, 3, 3)
In [484]: np.allclose(c, f.transpose(1,2,0))
Out[484]: True
我认为5
维度是不断前进".这就是循环的作用.在einsum
中,i
的所有部分都相同.
I think of the 5
dimension as 'going-along-for-ride'. That's what your loop does. In einsum
the i
is the same in all parts.
这篇关于沿指定轴的两个3D矩阵之间的np.dot积的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!