3d 数组到矩阵乘法

Question

我有一个名为 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

Answer 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.

上找到关于这个非常强大的功能的更完整的解释以及如何使用它的许多示例。

Answer 2

@/matmul 很好地处理 batches，但规则是对于 3d 数组，第一个维度是批处理，dot 在最后 2 个维度上完成，与通常“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 维度，即结果的 'carried through'，以及共享的乘积之和j维度。

3d 数组到矩阵乘法

3d array to matrix multiplication

python

numpy

matrix

matrix-multiplication