python - 掩码矩阵乘法

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.

设 A 为:

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 reshape 为 Q x M，其中 Q=NxP。然后执行某种稀疏乘法，对于 bool 稀疏矩阵的每一行 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