对两个向量列表进行“逐元素”点积的 numpy.einsum 更快的替代方案?

3 投票
2 回答
1999 浏览
提问于 2025-04-18 09:54

假设你有两个向量数组:

v1 = np.array([ [1, 2], [3, 4] ])

v2 = np.array([ [10, 20], [30, 40]])

我们想要生成一个新的数组,它的内容相当于:

v3 = np.array([ np.dot(v1[0], v2[0]), np.dot(v1[1], v2[1]) ])

目前我使用的是:

v3 = np.einsum('ij,ij->i', v1, v2)

不过,我在代码中这样做的次数非常多,所以如果能加快速度对我来说会很有帮助。

我们怎么能加快这个过程呢?虽然np.einsum已经相当高效了,但我在想,对于这个特定的用例,有没有更快的方法呢?

2 个回答

1

试试 inner1d

import numpy as np
import cProfile
from numpy.core.umath_tests import inner1d

v1 = np.random.random((10**7,2,))  # 10 million random vectors
v2 = np.random.random((10**7,2,))  # 10 million random vectors
v3 = np.einsum('ij,ij->i', v1, v2) # einsum
v4 = inner1d(v1,v2)                # inner1d (~2x faster than einsum)

cProfile.run("np.einsum('ij,ij->i', v1, v2)") # cProfile: 3 function calls in 0.065 seconds
cProfile.run("inner1d(v1,v2)") # cProfile: 2 function calls in 0.033 seconds

print np.allclose(v3,v4) # Returns True
2

einsum 是我能想到的三种选择中最好的一个:

In [73]: timeit v3=np.einsum('ij,ij->i',v1,v2)
100000 loops, best of 3: 5.14 us per loop

In [74]: timeit np.diag(np.dot(v1,v2.T))
100000 loops, best of 3: 7.43 us per loop

In [75]: timeit np.sum(v1*v2,axis=1)
100000 loops, best of 3: 16.8 us per loop

有几个问题需要考虑:

  • 这个计算真的那么耗费资源吗?
  • 如果确实比较耗费资源,你需要经常进行这样的计算吗?
  • 能否把多个 einsum 的调用合并起来 - 把数组连接在一起?

撰写回答