如何在NumPy中有效地计算列向外向量

Question

假设矩阵A具有m x r大小，而B具有n x r。在NumPy中，如何在A和B之间进行列级外积？

也就是说，我想要一个矩阵，C (大小：m x n x r)，其中C[:, :, i]是A[:, i]和B[:, i]的外部乘积。

当然，我可以在列上使用for循环，但是我想知道是否有一种矢量化的方法来加快计算速度？

Answer 1

在这种情况下，我通常会调用einsum以便于矢量化：

>>> m,n,r = 30,50,70
>>> A = np.random.random((m,r))
>>> B = np.random.random((n,r))
>>> C = np.einsum('ik,jk->ijk', A, B)
>>> C2 = np.dstack([np.outer(A[:,i], B[:,i]) for i in range(r)])
>>> C.shape
(30, 50, 70)
>>> np.allclose(C, C2)
True

这样，我们就可以清楚地知道我们希望在每个坐标上发生什么，同时仍然比循环更快：

>>> %timeit C = np.einsum('ik,jk->ijk', A, B)
The slowest run took 4.83 times longer than the fastest. This could mean that an intermediate result is being cached.
1000 loops, best of 3: 216 µs per loop

>>> %timeit C2 = np.dstack([np.outer(A[:,i], B[:,i]) for i in range(r)])
100 loops, best of 3: 2.08 ms per loop