使用 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])

我从链接 post 改编 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 的 last 轴。所以你可以想像做每个内积一样np.einsum('jk,k->j', a[0], b[0])。但是我们正在为 a 和 b 的每个元素做一个，因此包含第一个轴，它是 而不是 收缩的。因此，ijk,ik->ij.

使用 numpy 高效计算多个点积

Efficiently compute multiple dot products with numpy

python

numpy

linear-algebra