屏蔽矩阵乘法

Question

我有一个A = N x M矩阵和另一个B = N x P x M数组，其中p通常是9或15。对于来自A的每个向量a，必须将其与来自同一行的每个B的pi相乘，以获得维度N x P的输出

我正在使用numpy和Python，并将在GPU上执行此操作

举个小例子，让N=4，M=5，P=3

让我们来谈谈：

array([[0.18503431, 0.2628188 , 0.26343728, 0.8356702 , 0.47581551],
       [0.70827725, 0.04006919, 0.58975722, 0.90874113, 0.43946412],
       [0.40669507, 0.63328008, 0.95832881, 0.59041436, 0.63578578],
       [0.12129919, 0.74470057, 0.62271405, 0.97760796, 0.6499647 ]])

让B成为：

array([[[4.29031165e-01, 6.17324572e-01, 6.54726975e-02, 1.72218768e-02, 3.53970827e-01],
        [3.38821841e-01, 3.80128792e-01, 7.70995505e-01, 7.38437494e-03, 5.87395036e-02],
        [4.75661932e-01, 3.75617802e-01, 1.28564731e-01, 3.66302247e-01, 6.70953890e-01]],

       [[8.96228996e-02, 1.67135584e-02, 4.56921778e-01, 8.25731354e-01, 7.66242539e-01],
        [5.16651815e-01, 4.27179773e-01, 9.34673912e-01, 2.04687170e-01, 7.68417953e-01],
        [5.90980849e-01, 5.03013376e-01, 8.41765736e-02, 8.08221224e-01, 7.76765422e-01]],

       [[3.25802668e-01, 8.58148960e-01, 9.47505735e-01, 1.01405305e-01, 8.34114717e-01],
        [1.65308159e-01, 9.74572631e-01, 2.69886016e-01, 7.44036253e-02, 4.73350521e-01],
        [8.59030672e-01, 3.96972621e-01, 7.34687493e-01, 2.84647032e-02, 7.19723378e-01]],

       [[1.35751242e-01, 1.74882898e-01, 5.48875709e-01, 7.33443675e-01, 4.05282650e-01],
        [8.41298770e-01, 6.24323279e-01, 5.83482185e-01, 4.28514313e-01, 1.96797205e-01],
        [7.93345700e-04, 3.01441721e-01, 7.59451146e-01, 9.09102382e-01, 7.11518948e-01]]])

我希望我的输出是这样的：

[[np.dot(a[0], b[0][0]), np.dot(a[0], b[0][1]), np.dot(a[0], b[0][2])],
[np.dot(a[1], b[1][0]), np.dot(a[1], b[1][1]), np.dot(a[1], b[1][2])],
[np.dot(a[2], b[2][0]), np.dot(a[2], b[2][1]), np.dot(a[2], b[2][2])],
[np.dot(a[3], b[3][0]), np.dot(a[3], b[3][1]), np.dot(a[3], b[3][2])]]

手动执行此操作将提供：

[[0.44169455751462816, 0.3998276862221848, 0.845960080871557],
[1.4207326179275017, 1.4579799277670968, 1.564201768913105],
[2.174162453912622, 1.287925491552765, 1.779226448174152],
[1.4689343122491012, 1.4771555510001255, 2.0487088726424365]]

因为我想在GPU上实现这一点，这显然需要将我的问题转化为矩阵乘法（如果我不使用GPU，这是真的）。 但我不知道如何把它转换成那样

我的一个想法是将B重塑为Q x M，其中Q=NxP。然后执行某种稀疏乘法，其中对于布尔稀疏矩阵的每一行i，我打开第(0:P) + P*i个元素。（把它画出来是有意义的），但是我确实觉得有一种更优雅的方法来实现这一点，因为创建稀疏矩阵和执行操作可能需要时间，而且我的矩阵的稀疏性根本不是随机的

我怎样才能优雅地解决这个问题

请注意，我无法执行一些操作，例如广播/重复我的A矩阵P次，并执行巨大的矩阵乘法和挑选相关值，因为通常N和M会非常大（2000ish和256），但P会非常小，因此，对所有向量执行全局矩阵乘法意味着我将执行>；95%不必要的计算

Answer 1

您可以在这里使用einsum来有效地计算这个值

---

np.einsum('ij,ikj->ik', A, B) # or torch.einsum

---

>>> np.allclose(np.einsum('ij,ikj->ik', A, B), manual)
True

Answer 2

(np.expand_dims(A, 1) * B).sum(axis=2)

有点像“手动”矩阵乘法。首先，A被重塑为N x 1 x M。然后，它是按元素乘以B。Numpy将广播A形状以匹配B，因此最后我们将得到N x P x M。每一行将包含A[i] * B[i]，剩下要做的就是沿最后一个轴求和，得到点积