快速 Numpy 循环

Posted 2023-02-23

【中文标题】快速 Numpy 循环

【英文标题】：Fast Numpy Loops

【发布时间】：2016-10-14 02:03:15

【问题描述】：

您如何优化此代码（没有矢量化，因为这会导致使用计算的语义，这通常远非不重要） ：

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，它的性能非常好。

【问题讨论】：

cython 和 jit 都没有加速，因为您已经在运行 C 代码（通过 np.outer）。这里的问题实际上是循环本身，你需要改变它的内部结构，这样这些方法才能真正被加速。

我知道向量化内部（或两者）循环将显着加速代码。我的观点是，显然循环会产生一些不应该存在的显着开销。换句话说：为什么调用 np.outer 比在具有 200 行（矢量化）的矩阵上调用 np.outer 慢 200 倍，而不是说 Matlab 循环，这不是问题......能克服吗？

我认为我不能再提供任何帮助，但请看一下这个关于每个（Python 和 Matlab）如何处理循环的答案：***.com/a/17242928/2752305

嗯，一件事是调用它 200 次的函数开销。这在 Python 和 MATLAB 级别都会变慢。尽管最近 JIT 对其进行了显着改进，但 NumPy 可能需要赶上这一点（没有太多关于它的信息）。

另外，您没有拨打np.outer 200 次。您正在调用它 40000 次。

【参考方案1】：

这是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 必须能够用更快的代码替换“外部”——它会重写迭代。但我对 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() 的调用相比没有。

【参考方案2】：

如果内存允许，您可以使用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)

【讨论】：

感谢，但我已经知道向量化代码可以加快计算速度。有时很容易看出如何对代码进行矢量化（就像这里使用 einsum 所做的那样），但有时需要对潜在问题有非常深入的了解，并且在循环中编写代码要容易得多。那该怎么办呢？

@ndbd 如果您要求一个关于如何加速代码的通用案例，我会说这取决于。但是根据我的个人经验，我发现 NumPy ufunc 和函数（如einsum）和基于点积的函数在我们处理 Python 级别的向量化方法的乘法和减法时很有用。对于一般情况，我真的不能说任何我认为值得注意的事情，对不起！

【参考方案3】：

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

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

使用 Numba 或 Cython 不会优化这些问题！相反，这些工具允许您编写比普通 Python 快得多的循环代码。

另外，特别是对于 Numba，您应该注意 "object mode" and "nopython mode" 之间的区别。您示例中的紧密循环必须在 nopython 模式下运行才能提供任何显着的加速。但是，numpy.outer 是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

【参考方案4】：

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

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