我正在try 使用numba并行化一个简单的python程序.它由两个功能组成.第一种是简单幂法

@numba.jit(nopython=True)
def power_method(A, v):
    u = v.copy()
    for i in range(3 * 10**3):
        u = A @ u
        u /= np.linalg.norm(u)
    return u

第二个函数在网格上与向量v迭代,并为各种向量v运行power_method函数.

@numba.jit(nopython=True, parallel=True)
def iterate_grid(A, scale, sz):
    assert A.shape[0] == A.shape[1] == 3
    n = A.shape[0]

    results = np.empty((sz**3, n))
    tmp = np.linspace(-scale, scale, sz)
    
    for i1 in range(sz):
        v = np.empty(n, dtype=np.float64)
        v1 = tmp[i1]
        for i2, v2 in enumerate(tmp):
            for i3, v3 in enumerate(tmp):
                v[0] = v1
                v[1] = v2
                v[2] = v3

                u = power_method(A, v)

                idx = i1 * sz**2 + i2 * sz + i3
                results[idx] = u.copy()
    
    return results

然后我用

n = 3
A = np.random.randn(n, n)
iterate_grid(A, 5.0, 20)

所有迭代都是独立的.此外,理想情况下,计算会落在缓存中,因此我希望将第一个循环与prange并行将得到近似的线性加速度.

然而,顺序代码的壁时间为6.07秒,而并行代码的壁时间为6.07秒

@numba.jit(nopython=True, parallel=True)
def iterate_grid(A, scale, sz):
    assert A.shape[0] == A.shape[1] == 3
    n = A.shape[0]

    results = np.empty((sz**3, n))
    tmp = np.linspace(-scale, scale, sz)
    
    for i1 in numba.prange(sz):
        v = np.empty(n, dtype=np.float64)
        v1 = tmp[i1]
        for i2, v2 in enumerate(tmp):
            for i3, v3 in enumerate(tmp):
                v[0] = v1
                v[1] = v2
                v[2] = v3

                u = power_method(A, v)

                idx = i1 * sz**2 + i2 * sz + i3
                results[idx] = u.copy()
    
    return results

挂机时间为7.79秒.也就是说,在这种情况下,并行化会减慢代码的速度.

此外,正如我从iterate_grid.parallel_diagnostics(level=4)中看到的,numba熔断了

并行化power方法并不是我真正想要解决的任务,它只是一个小例子,numba的行为并没有达到我的预期.你能给我解释一下numba的这种行为,并建议我如何在网格上并行化迭代以获得线性加速度吗?

提前感谢您的帮助!

推荐答案

因此,在try 衡量实现的性能之前,您需要判断结果是否正确.结果表明,顺序实现和并行实现之间的结果是不同的(您可以使用np.allclose来判断或绘制数组差异).由于除了prange的使用之外,并行代码基本上是相同的,所以很可能是由于竞争条件造成的.多次运行程序并判断结果是否不同是判断可能存在的竞争条件的好方法(注意,没有差异并不意味着没有竞争条件,其他问题可能导致不确定的结果).在我的机器上进行的实际测试给出了非常不同的结果.

事情是your code looks completely correct!通常这意味着问题是由于未定义的行为或编译器/运行时错误造成的.更深入的分析表明,如果在并行循环中移动tmp = np.linspace(-scale, scale, sz),则不会出现问题.这当然是由于bug in Numba.

假设结果是正确的,那么性能很大程度上取决于power_method功能,这是无效的.事实上,Numba和Numpy都没有针对非常小的数组进行优化.Numba does not often optimize out array allocations除非在某些特定情况下(例如在基于prange的循环中使用检测到的不变量进行显式分配,或者如@max9111所指出的similar cases):数组是使用quite slow allocation方法在堆上分配的.除此之外,allocations do not scale可防止任何并行加速.优化是使用3个变量手动执行操作.以下是一个优化的实现:

@numba.njit
def fast_power_method(A, v1, v2, v3):
    # Unpacking
    # Note: there is no need for a copy here
    u1, u2, u3 = v1, v2, v3
    A11, A12, A13 = A[0]
    A21, A22, A23 = A[1]
    A31, A32, A33 = A[2]

    for i in range(3_000):
        # Optimized matrix multiplication
        t1 = u1 * A11 + u2 * A12 + u3 * A13
        t2 = u1 * A21 + u2 * A22 + u3 * A23
        t3 = u1 * A31 + u2 * A32 + u3 * A33

        # Renormalization
        # Note: multiplications are faster than divisions
        norm = np.sqrt(t1**2 + t2**2 + t3**2)
        inv_norm = 1.0 / norm
        u1 = t1 * inv_norm
        u2 = t2 * inv_norm
        u3 = t3 * inv_norm

    return u1, u2, u3

@numba.njit
def iterate_grid(A, scale, sz):
    assert A.shape[0] == A.shape[1] == 3
    n = A.shape[0]

    results = np.empty((sz**3, n))
    tmp = np.linspace(-scale, scale, sz)

    for i1 in range(sz):
        v = np.empty(n, dtype=np.float64)
        v1 = tmp[i1]
        for i2, v2 in enumerate(tmp):
            for i3, v3 in enumerate(tmp):
                u1, u2, u3 = fast_power_method(A, v1, v2, v3)

                idx = i1 * sz**2 + i2 * sz + i3
                results[idx, 0] = u1
                results[idx, 1] = u2
                results[idx, 2] = u3
    
    return results

此实现需要0.4秒,而初始版本大约需要10秒.因此,它是25 times faster,但仍然是连续的.由于浮点舍入和非关联性,结果略有不同.请注意,在我的6核机器上使用prange这个版本速度快5倍,所以它的扩展性很好,但由于同样的原因,结果仍然是错误的.这至少表明,分配无疑是瓶颈.

经验法则是check results,而micro-optimizations可以比使用多核(尤其是对于大量数字代码)更快地执行.

UPDATE:我填写了一个可用的Nuba bug here.

Python相关问答推荐

循环测试文件并为每个文件绘制

使用 Python 类型提示指示来自副作用的变量值

threading.get_ident() 在运行 pytest 时在不同线程之间返回相同的 ID

如何使枚举不计算空白索引?

如何使用 groupby 判断是否有多个版本

Python - 为什么在每次完整的 for 循环迭代后 y 的值都会增加?

FunctionTransformer & 在管道中创建新列

通过 API 获得的 Youtube 视频“重播次数”数据

为什么大小为 2 的列表不适用于这个嵌套的 for 循环

我无法使用 ctypes cuda 获得输出数字

如何从 Pandas 数据框中的一行中 Select 具有最高值的 3 列?

有什么方法可以获取公会的通知设置吗?

从第 1 列搜索共同值(患者 ID),如果其他列(病理)中的所有值都为空,则删除这些共同 ID 的行

如何使用 PyPDF2 更改已更改 PDF 的保存位置

奇怪的 datetime.utcnow() 错误

使用 Python 读取 dbus 属性

Matplotlib 将动画保存为视频,但内容为空

使用ffmpeg python阅读视频时如何忽略自动旋转?

Sklearn 管道转换特定列 - ValueError:要解包的值太多(预期为 2)

使用 Selenium 提取/保存元素的文本和图像