沿指定轴的两个 3D 矩阵之间的 np.dot 积答案

【问题标题】：np.dot product between two 3D matrices along specified axis沿指定轴的两个 3D 矩阵之间的 np.dot 积
【发布时间】：2018-08-23 04:50:49
【问题描述】：

我有两个 3D 矩阵：

a = np.random.normal(size=[3,2,5])
b = np.random.normal(size=[5,2,3])

我想要每个切片分别沿 2 轴和 0 轴的点积：

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 来做到这一点（为了效率和速度）我试过了：

c = np.tensordot(a, b, axes=[2,0])

但我得到一个包含 36 个元素（而不是 45 个）的 4D 数组。 c.shape, c.size = ((3L, 2L, 2L, 3L), 36)。我在这里发现了一个类似的问题（Numpy tensor: Tensordot over frontal slices of tensor），但这并不是我想要的，我无法将该解决方案推断为我的问题。总结一下，我可以使用 np.tensordot 来计算上面显示的 c 数组吗？

更新 #1

@hpaulj 的答案是我想要的，但是在我的系统（python 2.7 和 np 1.13.3）中，这些方法非常缓慢：

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))

【问题讨论】：

标签： python numpy matrix-multiplication tensor numpy-ndarray

【解决方案1】：

使用einsum 比使用tensordot 更容易。所以让我们从这里开始：

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）。

在测试想法时，如果c 的前两个维度不同，可能会有所帮助。把它们混在一起的机会就更少了。

在tensordot 中指定dot 求和轴（轴）很容易，但很难约束其他维度的处理。这就是你得到一个 4d 数组的原因。

我可以让它与转置一起工作，然后是对角线：

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

我计算的值比需要的多得多，并且大部分都扔掉了。

matmul 加上一些转置可能会更好。

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 在所有部分中都是相同的。

【讨论】：