在 Python 中使用数组更快的 for 循环

Posted 2023-02-24

【中文标题】在 Python 中使用数组更快的 for 循环

【英文标题】：Faster for-loops with arrays in Python

【发布时间】：2021-02-08 22:20:35

【问题描述】：

N, M = 1000, 4000000
a = np.random.uniform(0, 1, (N, M))
k = np.random.randint(0, N, (N, M))

out = np.zeros((N, M))
for i in range(N):
    for j in range(M):
        out[k[i, j], j] += a[i, j]

我使用很长的 for 循环；上面的%%timeit 用pass 替换操作产量

1min 19s ± 663 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

这在上下文中是不可接受的（C++ 花了 6.5 秒）。没有理由使用 Python 对象完成上述操作；数组具有明确定义的类型。在 C/C++ 中将其作为扩展来实现对开发人员和用户来说都是一种过度杀伤。我只是将数组传递给循环并进行算术运算。

有没有办法告诉 Numpy“将此逻辑移至 C”，或者另一个可以处理仅涉及数组的嵌套循环的库？我为一般情况寻找它，而不是针对这个特定示例的解决方法（但如果你有一个，我可以打开一个单独的问答）。

【问题讨论】：

您是否正在寻找用于准 Python 的编译器，如 Cython、numba 等或其他方法（如矢量化 numpy 操作）？

@MisterMiyagi 不确定您所说的“编译器”是什么意思，任何拥有 Python 标准 C 实现的人都应该能够运行它，并且额外的库应该处理它们自己的依赖项（如果需要，包括编译器）。 -- 并且没有向量化等，这不是算法优化问题。

@MisterMiyagi：Cython 似乎因为不想要自己的扩展而被排除在外，但 numba 是一种可能性，假设转换为矢量化 numpy 操作是不可行的。

@OverLordGoldDragon：对于类似的事情，@numba.jit(nopython=True) would be the first thing I'd think of。我不能说它是否会完全优化您的案例，但值得一试（这是迄今为止最简单的调整）。我会注意到，您呈现的代码不在函数中，这会使标准 CPython 变慢（只需将其包装在函数中就会将变量的每次读/写从 dict 查找更改为 C 数组索引操作） .

@OverLordGoldDragon：我在本地机器上测试过。不包装函数（但将M 减少到 40000），大约需要 29.1 秒的用户时间；将其包装在一个函数中使其下降到 25.5 秒（小但有意义的变化），并用 @numba.jit(nopython=True) 装饰该函数将其下降到 2.5 秒（尽管它第一次运行挂钟时间约为 12.4 秒，第二次运行下降到 3.6；加载 numba 本身和 jiting 有一些不小的启动成本，特别是在我的情况下，库必须第一次从 NFS 缓存）。

【参考方案1】：

这基本上是Numba 背后的想法。 不如 C 快，但它可以接近...它使用 jit 编译器将 python 代码编译到机器上，并且与大多数 Numpy 函数兼容。 （在文档中您可以找到所有详细信息）

import numpy as np
from numba import njit


@njit
def f(N, M):
    a = np.random.uniform(0, 1, (N, M))
    k = np.random.randint(0, N, (N, M))

    out = np.zeros((N, M))
    for i in range(N):
        for j in range(M):
            out[k[i, j], j] += a[i, j]
    return out


def f_python(N, M):
    a = np.random.uniform(0, 1, (N, M))
    k = np.random.randint(0, N, (N, M))

    out = np.zeros((N, M))
    for i in range(N):
        for j in range(M):
            out[k[i, j], j] += a[i, j]
    return out

纯 Python：

%%timeit

N, M = 100, 4000
f_python(M, N)

每个循环 338 毫秒 ± 12.6 毫秒（平均值 ± 标准偏差，7 次运行，每个循环 1 个）

使用 Numba：

%%timeit

N, M = 100, 4000
f(M, N)

每个循环 12 毫秒 ± 534 微秒（平均值 ± 标准偏差。7 次运行，每次 100 次循环）

【讨论】：

优秀。另外，您的示例支持 Python 的 for 循环开销的 smashing。我会把这个问题留得更久一点，但我怀疑有什么能胜过这个。

@OverLordGoldDragon 你真的相信它在 451 ns 内完成了 4 亿次迭代吗？

@HeapOverflow 不久之后我意识到这是荒谬的，为了测试我将一个带有 +=1 的数组传递到循环中，以确保循环没有被跳过，并且执行时间几乎没有变化 - 然后可能已经完成了其他一些优化，但我没有进一步追求它。 -- 对于其他人，堆指的是所需的 4e8/5e-7 = 1e15 Hz CPU（一个核心）。