python - 在 NumPy 中向量化多个向量矩阵乘法

Question

我无法弄清楚如何排列我的轴，以便我可以以矢量化方式执行以下操作。

本质上我有一个向量数组，一个矩阵数组，我想为每个对应的向量 V 和矩阵 M 评估 VMV^T

import numpy as np

N = 5 # normally 100k or so
vecs = np.random.rand(N, 2)
mats = np.random.rand(N, 2, 2)

output = np.array([np.dot(np.dot(vecs[i, ...], mats[i, ...]), vecs[i, ...].T) for i in range(N)])

如果它更简单，对下面的中间结果进行矢量化也会有帮助：

intermediate_result = np.array([np.dot(vecs[i, ...], mats[i, ...]) for i in range(N)])
# then I can do
output = np.sum(intermediate_result * vecs, axis=-1)

Answer 1

对于einsumN = 100k，基于 - 的解决方案比您的循环快 100 倍：

%timeit np.array([np.dot(np.dot(vecs[i, ...], mats[i, ...]), vecs[i, ...].T) for i in range(N)])
%timeit np.einsum('...i,...ij,...j->...', vecs, mats, vecs)
np.allclose(np.array([np.dot(np.dot(vecs[i, ...], mats[i, ...]), vecs[i, ...].T) for i in range(N)]),
            np.einsum('...i,...ij,...j->...', vecs, mats, vecs))
1 loops, best of 3: 640 ms per loop
100 loops, best of 3: 7.02 ms per loop
Out[45]: True