我有两个形状为 (N, 2, 2) 的 3d 数组 A 和 B，我想根据 N 轴将每个元素与每个 2x2 矩阵上的矩阵乘积相乘。使用循环实现，它看起来像
C[i] = dot(A[i], B[i])
有没有办法在不使用循环的情况下做到这一点？我研究了 tensordot，但无法让它工作。我想我可能想要tensordot(a, b, axes=([1,2], [2,1])) 之类的东西，但这给了我一个 NxN 矩阵。
【问题讨论】：
标签： python arrays numpy matrix vectorization
您似乎正在对沿第一个轴的每个切片进行矩阵乘法。同样，您可以像这样使用np.einsum -
np.einsum('ijk,ikl->ijl',A,B)
我们也可以使用np.matmul -
np.matmul(A,B)
在 Python 3.x 上，此 matmul 操作使用 @ operator 进行了简化 -
A @ B
方法-
def einsum_based(A,B):
return np.einsum('ijk,ikl->ijl',A,B)
def matmul_based(A,B):
return np.matmul(A,B)
def forloop(A,B):
N = A.shape[0]
C = np.zeros((N,2,2))
for i in range(N):
C[i] = np.dot(A[i], B[i])
return C
时间安排 -
In [44]: N = 10000
...: A = np.random.rand(N,2,2)
...: B = np.random.rand(N,2,2)
In [45]: %timeit einsum_based(A,B)
...: %timeit matmul_based(A,B)
...: %timeit forloop(A,B)
100 loops, best of 3: 3.08 ms per loop
100 loops, best of 3: 3.04 ms per loop
100 loops, best of 3: 10.9 ms per loop
【讨论】：
einsum 的工作原理。你读过哪些东西可以教你语法是如何工作的？
permute 所做的那样，随便玩弄它。最近，又想通了how to use np.einsum to replace np.any。
np.einsum 支持的 NumPy 了吗？在这种情况下，我不确定如何让 tensordot 工作。
你只需要在你的张量的第一维上执行操作，它由0标记：
c = tensordot(a, b, axes=(0,0))
这将如您所愿。此外，您不需要轴列表，因为它只是沿着您正在执行操作的一个维度。使用axes([1,2],[2,1])，您可以交叉乘以第 2 维和第 3 维。如果你用索引符号（爱因斯坦求和约定）写它，这对应于c[i,j] = a[i,k,l]*b[j,k,l]，因此你正在收缩你想要保留的索引。
编辑：好的，问题是两个 3d 对象的张量积是 6d 对象。由于收缩涉及成对的索引，因此您无法通过 tensordot 操作获得 3d 对象。诀窍是将您的计算分成两部分：首先您在索引上执行tensordot 以执行矩阵运算，然后您采用张量对角线以将您的 4d 对象减少到 3d。在一个命令中：
d = np.diagonal(np.tensordot(a,b,axes=()), axis1=0, axis2=2)
张量表示法d[i,j,k] = c[i,j,i,k] = a[i,j,l]*b[i,l,k]。
【讨论】：
(2, 2, 2, 2)，而不是(10, 2, 2)，所以可能有问题。