Python 用tensordot实现批量矩阵乘法

我试图实现与np.matmul并行矩阵乘法相同的行为，只使用张量、点和整形等

---

我将此转换为使用的库没有支持并行乘法的matmul，只有点和张量

此外，我希望避免在第一维度上进行迭代，并希望使用一组矩阵乘法和重塑来实现这一点（希望使用BLAS/GPU来运行它，因为我有大量的小矩阵要并行计算）

以下是一个例子：

import numpy as np

angles = np.array([np.pi/4, 2*np.pi/4, 2*np.pi/4])

vectors = np.array([ [1,0],[1,-1],[-1,0]])

s = np.sin(angles)
c = np.cos(angles)

rotations = np.array([[c,s],[-s,c]]).T

print rotations

print vectors

print("Correct: %s" % np.matmul(rotations, vectors.reshape(3,2,1)))

# I want to do this using tensordot/reshaping, i.e just gemm BLAS operations underneath
print("Wrong: %s" % np.tensordot(rotations, vectors, axes=(1,1)))

其输出为：

Correct: [[[  7.07106781e-01]
  [  7.07106781e-01]]

 [[  1.00000000e+00]
  [  1.00000000e+00]]

 [[ -6.12323400e-17]
  [ -1.00000000e+00]]]


Wrong: [[[  7.07106781e-01   1.11022302e-16  -7.07106781e-01]
  [ -7.07106781e-01  -1.41421356e+00   7.07106781e-01]]

 [[  6.12323400e-17  -1.00000000e+00  -6.12323400e-17]
  [ -1.00000000e+00  -1.00000000e+00   1.00000000e+00]]

 [[  6.12323400e-17  -1.00000000e+00  -6.12323400e-17]
  [ -1.00000000e+00  -1.00000000e+00   1.00000000e+00]]]

是否有一种方法可以修改第二个表达式以获得与第一个表达式相同的结果，只需使用dot/tensordot

---

我相信这是可能的，并且已经看到了，但从来没有任何例子

我们需要使一个保持一致，并在输出中保持一致。所以，

tensordot/dot

在这里不起作用。也许可以解释为什么不会。但是，我们可以使用，在大多数情况下（根据我的经验），它比

np.matmul

稍微快一些

实现看起来像这样-

np.einsum('ijk,ik->ij',rotations, vectors)

此外，所需的输出似乎有一个尾随单态dim。因此，在那里附加一个带有

None/np.newaxis

的新轴，如下所示-

np.einsum('ijk,ik->ij',rotations, vectors)[...,None]

---

“我正在使用的库没有支持并行乘法的matmul，只有点和张量。”-我也不能使用einsum。还有其他想法吗？@ChrisBamford这是什么图书馆？它是tensorflow吗？@ChrisBamford不，它不可能以矢量化的方式使用

dot/tensordot

。所以，我想说，只需循环。

tensordot

交换和重新塑造，使问题减少到一个

dot

（然后返回）。一些

matmul

操作可以通过从一个更大的“外部”计算中获取对角线来实现。是的，我确实注意到了获取对角线，但我认为这可能比循环效率低得多