Question

我有一个A = N x M矩阵和另一个数组B = N x P x M，其中P通常是9或15。对于a中的每个向量A，必须将其与同一行pi中的每个B相乘，以得到尺寸为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重塑为Q x M，其中Q=NxP。然后执行某种稀疏乘法，对于布尔稀疏矩阵的每一行i，我打开第(0:P) + P*i个元素。（将其绘制出来是有道理的），但是我当然感觉有一种更优雅的方法来执行此操作，因为创建稀疏矩阵并执行操作可能需要一些时间，并且矩阵的稀疏性根本不是随机的。 >

我该如何优雅地解决这个问题。

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

Answer 1

(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]，剩下要做的就是沿最后一个轴求和以获得点积。

Answer 2

您可以在此处使用einsum来进行有效计算。

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

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

屏蔽矩阵乘法

2 个答案: