python 中的高效外积

Efficient outer product in python

当我们必须处理 10k 维向量时,python 中的外积似乎很慢。有人可以告诉我如何在 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

由于我必须多次执行此操作,所以我的代码越来越慢。

它应该像使用 numpy.outer() 一样简单:将在 C 中实现高性能的单个函数调用。

没有比这更快的了,这些是您的选择:

numpy.outer

>>> %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

鹦鹉

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

赛通

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

@elyase 的回答很好,并且被正确地接受了。这是一个额外的建议,如果您可以使用它,可能会更快地调用 np.outer

你说"I have to do this operation several times",所以有可能你可以重用保存外积的数组,而不是每次都分配一个新的。这可以大大提高性能。

首先,使用一些随机数据:

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))

现在使用out作为np.outer的第三个参数:

In [36]: %timeit np.outer(a, b, out)
100 loops, best of 3: 2.38 ms per loop

因此,如果您可以重用包含外积的数组,您将获得不错的性能提升。

如果您使用 einsumout 参数,并且在 cython 函数中,如果您为输出添加第三个参数而不是在带有 np.empty。 (@elyase 答案中的其他 compiled/jitted 代码也可能会从中受益,但我只尝试了 cython 版本。)

注意事项! 以上显示的好处在实践中可能无法实现。 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