使用 3d 数组的高效 numpy 行矩阵乘法答案

【问题标题】：Efficient numpy row-wise matrix multiplication using 3d arrays使用 3d 数组的高效 numpy 行矩阵乘法
【发布时间】：2022-09-22 21:11:52
【问题描述】：

我有两个形状为 (N, M, D) 的 3d 数组，我想执行有效的逐行（在 N 上）矩阵乘法，以使结果数组的形状为 (N, D, D)。

显示我尝试实现的目标的低效代码示例如下所示：

N = 100
M = 10
D = 50
arr1 = np.random.normal(size=(N, M, D))
arr2 = np.random.normal(size=(N, M, D))
result = []
for i in range(N):
    result.append(arr1[i].T @ arr2[i])
result = np.array(result)

但是，由于循环，此应用程序对于大 N 非常慢。有没有更有效的方法来实现这种计算而不使用循环？我已经尝试通过 tensordot 和 einsum 找到解决方案，但无济于事。

当您说慢时，您是指给定示例值 100、10 和 50 的慢吗？

标签： python numpy

【解决方案1】：

向量化解决方案是交换arr1的最后两个轴：

>>> N, M, D = 2, 3, 4
>>> np.random.seed(0)
>>> arr1 = np.random.normal(size=(N, M, D))
>>> arr2 = np.random.normal(size=(N, M, D))
>>> arr1.transpose(0, 2, 1) @ arr2
array([[[ 6.95815626,  0.38299107,  0.40600482,  0.35990016],
        [-0.95421604, -2.83125879, -0.2759683 , -0.38027618],
        [ 3.54989101, -0.31274318,  0.14188485,  0.19860495],
        [ 3.56319723, -6.36209602, -0.42687188, -0.24932248]],

       [[ 0.67081341, -0.08816343,  0.35430089,  0.69962394],
        [ 0.0316968 ,  0.15129449, -0.51592291,  0.07118177],
        [-0.22274906, -0.28955683, -1.78905988,  1.1486345 ],
        [ 1.68432706,  1.93915798,  2.25785798, -2.34404577]]])

【讨论】：