3d 数组到矩阵乘法答案

【问题标题】：3d array to matrix multiplication3d 数组到矩阵乘法
【发布时间】：2022-01-10 15:57:12
【问题描述】：

我有一个名为 vec 的矩阵，它有两列，vec[:,0] 和 vec[:,1]。 P 包含两个矩阵，P[0,:,:] 和 P[1,:,:]。我想将 P[0,:,:] 与 vec 的第一列相乘，并将 P[1,:,:] 与 vec 的第二列相乘。但是，操作 P@vec 还给了我 P[0,:,:] 与 vec 的第二列的矩阵乘积以及 P[1,:,:] 与 vec 的第一列的矩阵乘积，这会减慢我的代码。

是否可以直接计算第 1 列到矩阵 1 和第 2 列到矩阵 2 的对而无需“关闭”产品？

import numpy as np

P=np.arange(50).reshape(2, 5, 5)
vec=np.arange(10).reshape(5,2)
have=P@vec
want=np.column_stack((have[0,:,0],have[1,:,1]))

have,want

【问题讨论】：

标签： python numpy matrix matrix-multiplication

【解决方案1】：

numpy 中有一个非常强大的函数，叫做np.einsum。它可以执行各种张量收缩、轴重新排序和矩阵乘法。对于您的示例，您可以使用

res = np.einsum('nij,jn->in', P, vec)

之后res 与want 完全相同。

这是如何工作的：你给np.einsum函数你的数组以及一个告诉函数如何乘以数组的签名（那个'nij,jn->in'字符串）。简而言之，您希望P 张量的第三个轴与vec 的第一个轴收缩。因此，您在签名字符串中选择相同的索引j，并将其留在-> 之后的部分中。如果索引出现在-> 的左侧和右侧，则仅进行广播，此处针对n 和i 索引进行。

可以在the corresponding numpy documentation 找到对这个非常强大的功能的更完整解释以及如何使用它的许多示例。

【讨论】：

【解决方案2】：

@/matmul 很好地处理batches，但规则是对于 3d 数组，第一个维度是批处理，dot 在最后两个维度上完成，通常的“A 的最后一个维度是第二个维度”到 B" 配对的最后一个。

花了一些时间来解读您的描述，但似乎您希望p 中的第一个成为批处理，vec 的最后一个成为批处理。这意味着vec 需要转换为 (2,5,1) 才能与 (2,5,5) p 一起使用。

In [176]: P@vec.T[:,:,None]
Out[176]: 
array([[[  60],
        [ 160],
        [ 260],
        [ 360],
        [ 460]],

       [[ 695],
        [ 820],
        [ 945],
        [1070],
        [1195]]])

结果是 (2,5,1)。我们可以挤出最后一个得到 (2,5)，但显然你想要一个 (5,2)

In [179]: (P@vec.T[:,:,None])[...,0].T
Out[179]: 
array([[  60,  695],
       [ 160,  820],
       [ 260,  945],
       [ 360, 1070],
       [ 460, 1195]])

np.einsum('nij,jn->in', P, vec) 的效果相同，n 作为 batch 维度被“传递”到结果，而乘积和则在共享的 j 维度上。

【讨论】：