时间戳

首页 > 解决方案 > 并行化 numba 循环

python - 并行化 numba 循环

问题描述

之前，我问了一个关于 Numba 无法并行化的相对简单循环的问题。结果证明，一个解决方案使所有循环都变得明确。

现在，我需要执行相同任务的更简单版本：我现在有数组alpha和beta分别具有形状(m,n)，(b,m,n),并且我想计算计算参数的 2D 切片的 Frobenius 乘积，并找到使该乘积最大化的 beta 切片。以前，有一个额外的、大的第一维度，alpha所以我并行化了这个维度；现在我想在第一个维度上并行化，beta因为当 b>1000 时计算变得昂贵。

如果我天真地修改适用于上一个问题的代码，我会得到：

@njit(parallel=True)
def parallel_value_numba(alpha,beta):
    dot = np.zeros(beta.shape[0])
        for i in prange(beta.shape[0]):
            for j in prange(beta.shape[1]):
                for k in prange(beta.shape[2]):
                    dot[i] += alpha[j,k]*beta[i, j, k]
        index=np.argmax(dot)
        value=dot[index]
    return value,index

但 Numba 出于某种原因不喜欢这样并抱怨道：

numba.core.errors.LoweringError: Failed in nopython mode pipeline (step: nopython mode backend)
scalar type memoryview(float64, 2d, C) given for non scalar argument #3

所以相反，我尝试了

@njit(parallel=True)
def parallel_value_numba_2(alpha,beta):
    product=np.multiply(alpha,beta)
    dot1=np.sum(product,axis=2)
    dot2=np.sum(dot1,axis=1)
    index=np.argmax(dot2)
    value=dot2[index]
    return value,index

只要您在将其传递给函数之前进行广播，它就会编译alpha，beta.shape并且原则上 Numba 能够并行化 numpy 操作。但它运行得非常缓慢，比串行的纯 Python 代码慢得多

def einsum_value(alpha,beta):
    dot=np.einsum('kl,jkl->j',alpha,beta)
    index=np.argmax(dot)
    value=dot[index]
    return value,index

所以，我当前的工作代码使用了最后一个实现，但是这个函数仍然是运行时的瓶颈，我想加快它的速度。谁能说服 Numba 以可观的速度并行化此功能？

标签： pythonnumpyparallel-processingnumba

解决方案

这并不完全是解决方案的答案，但格式化评论更难。

Numba 根据传递给函数的参数生成不同的代码。例如，您的代码适用于以下示例：

>>> alpha = np.random.random((5, 4))
>>> beta = np.random.random((3, 5, 4))
>>> parallel_value_numba(alpha, beta)
(5.89447648574048, 0)

为了诊断问题，有必要有一个导致问题的特定参数值的示例。

阅读错误消息，您似乎正在传递一个memoryview对象，但 Numba 可能不完全支持它。

作为旁注，您不需要prange在每个循环中都使用。只要预期的迭代次数大于机器中的内核数，通常在外循环中使用它就足够了。

推荐阅读