矩阵乘法，同时对矩阵中的元素进行子集化并存储在新矩阵中

Question

我正在尝试使用 as 变量调用numpy.matmul

• 维度(p, t, q)的矩阵A
• 维度为(r, t)的矩阵B
• 形状为r和p类别的categories向量，用于获取B的切片并定义A确实使用的索引。

乘法是使用每个类别的索引迭代完成的。 对于每个类别p_i ，我从 A 中提取一个子矩阵 ( t, q )。 然后，我将它们与 B ( x, t ) 的子集相乘，其中x是由r == p_i定义的掩码。 最后， (x, t)和(t, q)的矩阵乘法产生存储在S[x]中的输出(x, q) 。

我注意到我无法找出该算法的非迭代版本。 第一个片段描述了一个迭代解决方案。 第二个是对我希望得到的东西的尝试，其中所有内容都被计算为一步，并且可能会更快。 但是，这是不正确的，因为矩阵 A 具有三个维度而不是两个。 也许在 NumPy 中无法通过一次调用来完成此操作，并且通常会寻找建议/想法进行尝试。

谢谢！

import numpy as np
p, q, r, t = 2, 9, 512, 4
# data initialization (random)
np.random.seed(500)
S = np.random.rand(r, q)
A = np.random.randint(0, 3, size=(p, t, q))
B = np.random.rand(r, t)
categories = np.random.randint(0, p, r)

print('iterative') # iterative
for i in range(p):
    # print(i)
    a = A[i, :, :]
    mask = categories == i
    b = B[mask]
    print(b.shape, a.shape, S[mask].shape,
          np.matmul(b, a).shape)
    S[mask] = np.matmul(b, a)

print(S.shape)

一个简单的写下来的方法

S = np.random.rand(r, q)
print(A[:p,:,:].shape)
result = np.matmul(B, A[:p,:,:])

# iterative assignment
i = 0
S[categories == i] = result[i, categories == i, :]
i = 1
S[categories == i] = result[i, categories == i, :]

下一个片段将在乘法步骤中产生错误。

# attempt to multiply once, indexing all categories only once (not possible)
np.random.seed(500)
S = np.random.rand(r, q)
# attempt to use the categories vector
a = A[categories, :, :]
b = B[categories]
# due to the shapes of the arrays, this multiplication is not possible
print('\nsingle step (error due to shapes of the matrix a')
print(b.shape, a.shape, S[categories].shape)
S[categories] = np.matmul(b, a)
print(scores.shape)

iterative
(250, 4) (4, 9) (250, 9) (250, 9)
(262, 4) (4, 9) (262, 9) (262, 9)
(512, 9)


single step (error due to shapes of the 2nd matrix a).
(512, 4) (512, 4, 9) (512, 9)

Answer 1

In [63]: (np.ones((512,4))@np.ones((512,4,9))).shape
Out[63]: (512, 512, 9)

这是因为第一个数组被broadcasted到 (1,512,4)。 我想你想做的是：

In [64]: (np.ones((512,1,4))@np.ones((512,4,9))).shape
Out[64]: (512, 1, 9)

然后去掉中间的维度得到一个(512,9)。

其它的办法：

In [72]: np.einsum('ij,ijk->ik', np.ones((512,4)), np.ones((512,4,9))).shape
Out[72]: (512, 9)

Answer 2

要完全删除循环，你可以试试这个

bigmask = np.arange(p)[:, np.newaxis] == categories
C = np.matmul(B, A)
res = C[np.broadcast_to(bigmask[..., np.newaxis], C.shape)].reshape(r, q)

# `res` has the same rows as the iterative `S` but in the wrong order
# so we need to reorder the rows
sort_index = np.argsort(np.broadcast_to(np.arange(r), bigmask.shape)[bigmask])
assert np.allclose(S, res[sort_index])

虽然我不确定它是否比迭代版本快得多。