使用 numpy 高效计算多个点积

Question

我正在尝试有效地计算许多点积。这篇文章非常接近我想要做的事情，但我无法让它发挥作用。我有一个大的矩阵列表（a）和一个向量列表（b）。我想在他们之间做一系列的点积运算。这就是现在的工作：

import numpy as np

a # shape (15000,4,4)

b # shape (15000,4)

out = np.empty((15000,4))

for i in range(15000):

    out[i] = np.dot(a[i],b[i])

我从链接的帖子中调整 np.tensordot 或 np.einsum 的所有尝试都未能给我想要的东西。如果有人知道如何做到这一点，我将不胜感激。

Answer 1

爱因斯坦求和工作得很好：>>> a = np.random.randn(100, 4, 4)>>> b = np.random.randn(100, 4)>>> foo = np.einsum('ijk,ik->ij', a, b)>>> bar = np.zeros_like(foo)>>> for i, (ai, bi) in enumerate(zip(a, b)):        bar[i] = np.dot(ai, bi)>>> np.allclose(foo, bar)True为了稍微解释一下总和，请注意您正在收缩. b因此，您可以考虑将每个内积视为按np.einsum('jk,k->j', a[0], b[0]). 但是我们对aand的每个元素都做一个b，因此包含了第一个轴，它没有收缩。因此，ijk,ik->ij。