向量化NumPy中的多个向量 - 矩阵乘法
问题描述:
我无法确定如何安排我的坐标轴,以便我可以以矢量化的方式执行以下操作。向量化NumPy中的多个向量 - 矩阵乘法
基本上我有向量的阵列,矩阵的阵列,并且我要评估VMV^T为每个对应的矢量V和矩阵M
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)
答
的einsum
为基础的解决方案是比你的回路N = 100K速度100X:
%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
是的,'爱因斯坦'是惊人的:紧凑的记法,真快,灵活! – jorgeca 2013-05-13 09:29:55