python - 快速 Numpy 循环

Question

如何优化这段代码(没有向量化，因为这导致使用计算的语义，这通常远非重要) :

slow_lib.py:
import numpy as np

def foo():
    size = 200
    np.random.seed(1000031212)
    bar = np.random.rand(size, size)
    moo = np.zeros((size,size), dtype = np.float)
    for i in range(0,size):
        for j in range(0,size):
            val = bar[j]
            moo += np.outer(val, val)

要点是，此类循环通常对应于对某些向量运算进行双重求和的运算。

这很慢:

>>t = timeit.timeit('foo()', 'from slow_lib import foo', number = 10)
>>print ("took: "+str(t))
took: 41.165681839

好的，那么让我们对它进行 cynothize 并添加类型注释，就像没有明天一样:

c_slow_lib.pyx:
import numpy as np
cimport numpy as np
import cython
@cython.boundscheck(False)
@cython.wraparound(False)

def foo():
    cdef int size = 200
    cdef int i,j
    np.random.seed(1000031212)
    cdef np.ndarray[np.double_t, ndim=2] bar = np.random.rand(size, size)
    cdef np.ndarray[np.double_t, ndim=2] moo = np.zeros((size,size), dtype = np.float)
    cdef np.ndarray[np.double_t, ndim=1] val
    for i in xrange(0,size):
        for j in xrange(0,size):
            val = bar[j]
            moo += np.outer(val, val)


>>t = timeit.timeit('foo()', 'from c_slow_lib import foo', number = 10)
>>print ("took: "+str(t))
took: 42.3104710579

...呃...什么？ Numba 助你一臂之力!

numba_slow_lib.py:
import numpy as np
from numba import jit

size = 200
np.random.seed(1000031212)

bar = np.random.rand(size, size)

@jit
def foo():
    bar = np.random.rand(size, size)
    moo = np.zeros((size,size), dtype = np.float)
    for i in range(0,size):
        for j in range(0,size):
            val = bar[j]
            moo += np.outer(val, val)

>>t = timeit.timeit('foo()', 'from numba_slow_lib import foo', number = 10)
>>print("took: "+str(t))
took: 40.7327859402

那么真的没有办法加快速度吗？重点是:

• 如果我将内部循环转换为矢量化版本(构建一个更大的矩阵来表示内部循环，然后在更大的矩阵上调用 np.outer)，我会得到更快的代码。
• 如果我在 Matlab (R2016a) 中实现类似的东西，由于 JIT，这表现得相当好。

Answer 1

这是外部的代码:

def outer(a, b, out=None):    
    a = asarray(a)
    b = asarray(b)
    return multiply(a.ravel()[:, newaxis], b.ravel()[newaxis,:], out)

所以每次调用 outer 都涉及到一些 python 调用。那些最终调用编译代码来执行乘法。但是每个都会产生与数组大小无关的开销。

所以对 outer 的 200 (200**2?) 次调用将产生所有这些开销，而对所有 200 行的 outer 的一次调用将产生一次开销集，随后通过一次快速编译操作。

cython 和 numba 不会编译或以其他方式绕过 outer 中的 Python 代码。他们所能做的就是简化您编写的迭代代码 - 而这不会占用太多时间。

无需详细说明，MATLAB jit 必须能够用更快的代码替换“外部”——它会重写迭代。但我使用 MATLAB 的经验可以追溯到它出现之前的时间。

要真正提高 cython 和 numba 的速度，您需要一直使用原始的 numpy/python 代码。或者更好的办法是将精力集中在缓慢的内部部分。

用精简版本替换你的outer 将运行时间减少一半:

def foo1(N):
        size = N
        np.random.seed(1000031212)
        bar = np.random.rand(size, size)
        moo = np.zeros((size,size), dtype = np.float)
        for i in range(0,size):
                for j in range(0,size):
                        val = bar[j]
                        moo += val[:,None]*val   
        return moo

对于完整的 N=200，您的函数每个循环花费 17 秒。如果我用 pass 替换里面的两行(不计算)，每个循环的时间会下降到 3ms。换句话说，外部循环机制不是一个大的时间消耗者，至少与对 outer() 的许多调用相比是这样。