在 Numpy 中进行 2D 和 3D 矩阵乘法的 For 循环替代方案答案

【问题标题】：For-loop Alternative to do 2D & 3D Matrix Multiplication in Numpy在 Numpy 中进行 2D 和 3D 矩阵乘法的 For 循环替代方案
【发布时间】：2019-11-19 00:53:57
【问题描述】：

为了讨论，我有一个形状为 2x2 的 2D 矩阵 (A) 和一个形状为 2x2x10 的 3D 矩阵 (B)。我目前正在遍历矩阵 B 的最后一个轴，并一次构造一个完整的矩阵一个子矩阵。

A = np.random.random((2,2))
B = np.random.random((2,2,10))
C = np.zeros_like(B)
for i in range(B.shape[-1]):
    C[:,:,i] = A @ B[:, :, i]

实际上，我的矩阵比这大得多，我知道肯定有比 for 循环更有效的东西。我已经查看了几个先前的问题，其中解决方案涉及使用 np.tensordot 或 np.einsum，但坦率地说，我认为我没有正确使用它。

# Basic
C_basic = A @ B
print(f'Basic {np.allclose(C, C_basic)}')  # False

# Einsum
C_einsum = np.einsum('ij, jik-> ijk', A, B)
print(f'np.einsum {np.allclose(C, C_einsum)}')  # False

# Newaxis
C_newaxis = A[np.newaxis, ...] @ B
print(f'np.newaxis {np.allclose(C, C_newaxis)}')  # False

# Swapaxes
C_swapaxes = A @ np.swapaxes(B, 0, 2)
C_swapaxes = np.swapaxes(C_swapaxes, 0, 2)
print(f'np.swapaxes {np.allclose(C, C_swapaxes)}')  # False

【问题讨论】：

如果你的 B 是 (10,2,2)，那么 A@B 应该可以工作，结果是 (10,2,2)。如果 A 是 (2,3) 和 B 是 (3,4,10) 可能更容易测试。也就是说，使相加的轴与众不同（此处为 3）。在您的最后一种情况下，您需要另一个 swapaxes。在einsum
einsum 表达式应该是 'ij,jkl->ikl' - 点和在 j 上，其余点在 i 和 k 和 l作为批次（这是为了搭便车）。 @ 与 'i,j,k' 一起使用。批次维度必须是第一位的。

标签： python numpy matrix-multiplication

【解决方案1】：

这里有几种可能性：

import numpy as np

A = np.random.random((2,2))
B = np.random.random((2,2,10))
C = np.zeros_like(B)
for i in range(B.shape[-1]):
    C[:,:,i] = A @ B[:, :, i]

Cs = [np.einsum('ij,jkl',A,B),
      np.tensordot(A,B,((-1,),(0,))),
      (A@B.reshape(len(B),-1)).reshape(-1,*B.shape[1:]),
      np.moveaxis(A@np.moveaxis(B,-1,0),0,-1),
      (A@B.transpose(2,0,1)).transpose(1,2,0),
      np.inner(B.T,A).T,
      (B.T@A.T).T]

print([np.allclose(C,Ci) for Ci in Cs])

打印：

[True, True, True, True, True, True, True]

但是，这些并不是 100% 等效的：例如。前三个是 C 连续的，后两个 Fortran 和中间两个都不是。

您可以使用以下方法进行检查：

for Ci in Cs:
    print(Ci.flags)

【讨论】：