快速 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），我会得到 much更快的代码。
如果我在 Matlab (R2016a) 中实现类似的东西，由于 JIT，这表现得很好。

Answer 1

如果内存允许，您可以使用 np.einsum 以矢量化方式执行那些繁重的计算，就像这样 -

moo = size*np.einsum('ij,ik->jk',bar,bar)

也可以用np.tensordot-

moo = size*np.tensordot(bar,bar,axes=(0,0))

或者干脆 np.dot -

moo = size*bar.T.dot(bar)

Answer 2

这是 outer 的代码：

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 必须能够用更快的代码替换 'outer' - 它会重写迭代。但是我使用 MATLAB 的经验可以追溯到它的 jit 之前。

要真正提高 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().

的许多调用相比是这样。

Answer 3

Cython、Numba 等的许多教程和演示使这些工具看起来好像可以自动加速您的代码，但实际上，情况往往并非如此：您需要修改您的代码很少提取最佳性能。如果您已经实施了某种程度的矢量化，这通常意味着写出所有循环。 Numpy 数组操作不是最优的原因包括：

创建并循环了很多临时数组；
如果数组很小，每次调用的开销很大；
不能实现短路逻辑，因为数组是整体处理的；
有时无法使用数组表达式表达最佳算法，而您会选择时间复杂度较差的算法。

使用 Numba 或 Cython 不会优化这些问题！相反，这些工具允许您编写比普通 Python.

快得多的循环代码

此外，对于 Numba，您应该了解 "object mode" and "nopython mode". The tight loops from your example have to run in nopython mode to provide any significant speedup. However, numpy.outer is not yet supported by Numba 之间的区别，导致函数以对象模式编译。用jit(nopython=True)修饰，让这种情况抛出异常。

证明加速确实可行的示例：

import numpy as np
from numba import jit

@jit
def foo_nb(bar):
    size = bar.shape[0]
    moo = np.zeros((size, size))
    for i in range(0,size):
        for j in range(0,size):
            val = bar[j]
            moo += np.outer(val, val)
    return moo

@jit
def foo_nb2(bar):
    size = bar.shape[0]
    moo = np.zeros((size, size))
    for i in range(size):
        for j in range(size):
            for k in range(0,size):
                for l in range(0,size):
                    moo[k,l] += bar[j,k] * bar[j,l]
    return moo

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

np.allclose(foo_nb(bar), foo_nb2(bar))
# True

%timeit foo_nb(bar)
# 1 loop, best of 3: 816 ms per loop
%timeit foo_nb2(bar)
# 10 loops, best of 3: 176 ms per loop

Answer 4

您向我们展示的示例是一种低效算法，因为您多次计算相同的外积。最终的时间复杂度为 O(n^4)。可以减少到n^3.

for i in range(0,size):
    val = bar[i]
    moo += size * np.outer(val, val)

快速 Numpy 循环

Fast Numpy Loops

python

numpy

vectorization

cython