python - 使用tensordot实现批量矩阵乘法

Question

我正在尝试仅使用张量点、点和整形等实现与 np.matmul 并行矩阵乘法相同的行为。

我正在将其翻译成使用的库没有支持并行乘法的 matmul，只有 dot 和 tensordot。

此外，我想避免在第一个维度上进行迭代，并希望使用一组矩阵乘法和 reshape 来做到这一点(希望尽可能多地使用 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]]]

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

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

Answer 1

我们需要保持一个对齐并在输出端保持对齐。所以，tensordot/dot 在这里不起作用。 More info on tensordot可能会以某种方式解释为什么它不会。但是，我们可以使用 np.einsum ，在大多数情况下(根据我的经验)它被认为比 np.matmul 稍微快一些。

实现看起来像这样-

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

此外，所需的输出似乎有一个尾随的单例暗淡。所以，用 None/np.newaxis 附加一个新轴，就像这样 -

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