优雅的 Numpy 张量积

Question

我需要在 numpy（或 pytorch）中对两个张量进行乘积：

我有

A = np.arange(1024).reshape(8,1,128)
B = np.arange(9216).reshape(8, 128, 9)

并且想要获得C，点积在A (axis=2) 的最后一个暗淡和B (axis=1) 的中间暗淡上求和。这应该有尺寸8x9。目前，我正在做：

C = np.zeros([8, 9])
for i in range(8):
    C[i,:] = np.matmul(A[i,:,:], B[i,:,:])

如何优雅地做到这一点？

我试过了：

np.tensordot(weights, features, axes=(2,1)).

但它返回8x1x8x9。

Answer 1

一种方法是使用numpy.einsum。

C = np.einsum('ijk,ikl->il', A, B)

或者你可以使用broadcasted矩阵乘法。

C = (A @ B).squeeze(axis=1)
# equivalent: C = np.matmul(A, B).squeeze(axis=1)