For-loop另一种方法是在Numpy中进行2D和3D矩阵乘法

Question

为了讨论起见，我有一个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

Answer 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,))),
      ([email protected](len(B),-1)).reshape(-1,*B.shape[1:]),
      np.moveaxis([email protected](B,-1,0),0,-1),
      ([email protected](2,0,1)).transpose(1,2,0),
      np.inner(B.T,A).T,
      ([email protected]).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)