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

基于einsum的解决方案比N = 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

在NumPy中向量化乘法向量矩阵乘法

1 个答案: