在Python中,当我们需要处理10K级别维度的向量时,外积运算似乎非常缓慢。有没有人可以给我一些想法,如何在Python中加速这个操作?
代码如下:
In [8]: a.shape
Out[8]: (128,)
In [9]: b.shape
Out[9]: (32000,)
In [10]: %timeit np.outer(b,a)
100 loops, best of 3: 15.4 ms per loop
由于我必须多次执行此操作,我的代码变得越来越慢。
在Python中,当我们需要处理10K级别维度的向量时,外积运算似乎非常缓慢。有没有人可以给我一些想法,如何在Python中加速这个操作?
代码如下:
In [8]: a.shape
Out[8]: (128,)
In [9]: b.shape
Out[9]: (32000,)
In [10]: %timeit np.outer(b,a)
100 loops, best of 3: 15.4 ms per loop
由于我必须多次执行此操作,我的代码变得越来越慢。
>>> %timeit np.outer(a,b)
100 loops, best of 3: 9.79 ms per loop
numpy.einsum
>>> %timeit np.einsum('i,j->ij', a, b)
100 loops, best of 3: 16.6 ms per loop
numba
from numba.decorators import autojit
@autojit
def outer_numba(a, b):
m = a.shape[0]
n = b.shape[0]
result = np.empty((m, n), dtype=np.float)
for i in range(m):
for j in range(n):
result[i, j] = a[i]*b[j]
return result
>>> %timeit outer_numba(a,b)
100 loops, best of 3: 9.77 ms per loop
parakeet
from parakeet import jit
@jit
def outer_parakeet(a, b):
... same as numba
>>> %timeit outer_parakeet(a, b)
100 loops, best of 3: 11.6 ms per loop
cython
cimport numpy as np
import numpy as np
cimport cython
ctypedef np.float64_t DTYPE_t
@cython.boundscheck(False)
@cython.wraparound(False)
def outer_cython(np.ndarray[DTYPE_t, ndim=1] a, np.ndarray[DTYPE_t, ndim=1] b):
cdef int m = a.shape[0]
cdef int n = b.shape[0]
cdef np.ndarray[DTYPE_t, ndim=2] result = np.empty((m, n), dtype=np.float64)
for i in range(m):
for j in range(n):
result[i, j] = a[i]*b[j]
return result
>>> %timeit outer_cython(a, b)
100 loops, best of 3: 10.1 ms per loop
theano
from theano import tensor as T
from theano import function
x = T.vector()
y = T.vector()
outer_theano = function([x, y], T.outer(x, y))
>>> %timeit outer_theano(a, b)
100 loops, best of 3: 17.4 ms per loop
pypy
# Same code as the `outer_numba` function
>>> timeit.timeit("outer_pypy(a,b)", number=100, setup="import numpy as np;a = np.random.rand(128,);b = np.random.rand(32000,);from test import outer_pypy;outer_pypy(a,b)")*1000 / 100.0
16.36 # ms
╔═══════════╦═══════════╦═════════╗
║ method ║ time(ms)* ║ version ║
╠═══════════╬═══════════╬═════════╣
║ numba ║ 9.77 ║ 0.16.0 ║
║ np.outer ║ 9.79 ║ 1.9.1 ║
║ cython ║ 10.1 ║ 0.21.2 ║
║ parakeet ║ 11.6 ║ 0.23.2 ║
║ pypy ║ 16.36 ║ 2.4.0 ║
║ np.einsum ║ 16.6 ║ 1.9.1 ║
║ theano ║ 17.4 ║ 0.6.0 ║
╚═══════════╩═══════════╩═════════╝
* less time = faster
b [:,None] * a
。但它的时间与其他方法相差不大,在 outer
和 einsum
之间。相对排名会随着2个数组的大小而有所变化。 - hpauljfloat64
而不是 float
,否则它将无法编译。 - jubueche@elyase的回答很好,且已被正确接受。这里有一个额外的建议,如果您能使用它,可能会使调用np.outer
更快。
您说“我需要执行此操作多次”,因此可以重复使用保存外积的数组,而不是每次分配一个新数组。这可以提高性能。
首先,一些随机数据可供使用:
In [32]: a = np.random.randn(128)
In [33]: b = np.random.randn(32000)
以下是我电脑上np.outer(a, b)的基准时间:
In [34]: %timeit np.outer(a, b)
100 loops, best of 3: 5.52 ms per loop
out
数组来保存结果:In [35]: out = np.empty((128, 32000))
np.outer
的第三个参数改为out
:In [36]: %timeit np.outer(a, b, out)
100 loops, best of 3: 2.38 ms per loop
如果您可以重复使用保存外积的数组,则可以获得良好的性能提升。
如果您使用 einsum
的 out
参数,以及在cython函数中添加第三个参数作为输出,而不是使用 np.empty
在函数内进行分配,则可以获得类似的好处。(@elyase答案中的其他编译/加速代码也可能从中受益,但我只尝试了cython版本)
Nota bene!上面显示的收益在实践中可能无法实现。 out
数组适合我的CPU L3缓存,在由 timeit
命令执行的循环中使用时,它很可能仍然保留在缓存中。在实践中,该数组可能会在调用 np.outer
之间移出缓存。在这种情况下,改进并不那么显著,但仍然应该至少达到调用 np.empty()
的成本。
In [53]: %timeit np.empty((128, 32000))
1000 loops, best of 3: 1.29 ms per loop