python - 在 numpy 中广播 3D 数组操作的时间比例

Question

我正在尝试在两个 3D 数组上广播“>”的简单操作。一个具有维度 (m, 1, n)，另一个具有维度 (1, m, n)。如果我改变第三维 (n) 的值，我会天真地期望计算速度会随着 n 缩放。

但是，当我尝试明确地测量它时，我发现当 n 从 1 增加到 2 时，计算时间增加了大约 10 倍，之后缩放是线性的。

为什么从 n=1 到 n=2 时计算时间会急剧增加？我假设它是 numpy 中内存管理的产物，但我正在寻找更多细节。

下面附有代码和结果图。

import numpy as np
import time
import matplotlib.pyplot as plt

def compute_time(n):

    x, y = (np.random.uniform(size=(1, 1000, n)), 
            np.random.uniform(size=(1000, 1, n)))

    t = time.time()
    x > y 
    return time.time() - t

a = [
        [
            n, np.asarray([compute_time(n) 
            for _ in range(100)]).mean()
        ]
        for n in range(1, 30, 1)
    ]

a = np.asarray(a)
plt.plot(a[:, 0], a[:, 1])
plt.xlabel('n')
plt.ylabel('time(ms)')
plt.show()

广播操作的时间图

Answer 1

我无法证明这一点，但我很确定这是由于一种仅在 n==1 时可用的简单优化所致。

目前，numpy ufunc 实现基于计算机生成的最内层循环代码，映射到简单的 C 循环。封闭循环需要使用完全成熟的迭代器对象，这取决于负载，即最内层循环的大小和原子操作的成本可能是一个很大的开销。

现在，在 n==1 时，问题本质上是二维的(numpy 足够聪明，可以检测到)，最内层的循环大小为 1000，因此迭代器对象有 1000 步。从 n==2 向上，最内层循环的大小为 n，我们有 1,000,000 步的迭代器对象，这说明了您正在观察的跳跃。

正如我所说，我无法证明这一点，但我可以使它看起来合理:如果我们将变量维度移到前面，那么最内层的循环大小恒定为 1000，而外层循环在 1000 次迭代步骤中线性增长.事实上，这让跳跃消失了。

代码:

import numpy as np
import time
import matplotlib.pyplot as plt

def compute_time(n, axis=2):
    xs, ys = [1, 10], [10, 1]
    xs.insert(axis, n)
    ys.insert(axis, n)
    x, y = (np.random.uniform(size=xs),
            np.random.uniform(size=ys))

    t = time.perf_counter()
    x > y
    return time.perf_counter() - t

a = [
        [
            n,
            np.asarray([compute_time(n) for _ in range(100)]).mean(),
            np.asarray([compute_time(n, 0) for _ in range(100)]).mean()
        ]
        for n in range(0, 10, 1)
     ]

a = np.asarray(a)
plt.plot(a[:, 0], a[:, 1:])
plt.xlabel('n')
plt.ylabel('time(ms)')
plt.show()

相关:https://stackoverflow.com/a/48257213/7207392