时间戳

首页 > 解决方案 > 如何减少这个嵌套 for 循环的运行时间？

python - 如何减少这个嵌套 for 循环的运行时间？

问题描述

我正在编写一个将转换应用于一组 xyz 坐标的代码。旧的 xyz 坐标保存在数组 fea 中，新的 xyz 坐标分配给 newFea。npref 是较小的“原始”版本的 fea（也包含 xyz 坐标），而 W 是定义平移的矩阵。

这是代码：

@njit(cache=True, nogil=True)
def outerLoop(npref, fea, newFea, W):
n = len(W)
for i in range(1000):

    aTerms = W[n - 4] + fea[i][0] * W[n - 3] + fea[i][1] * W[n - 2] + fea[i][2] * W[n - 1]

    total = zeros((3,))

    for j in range(len(npref)):
        dist = sqrt(sum(square(npref[j] - fea[i])))
        u = square(dist) * sqrt(log(square(dist)))
        total += u * W[j]

    newFea[i] = total + aTerms
return newFea

目前，我正在调试它，我要去 1000，但实际上这将不得不运行到大约 130 万。我得到了大约 4 秒的运行时间，实际应用程序大约需要一个小时，我希望运行时间减少到 1 秒以下，以便完整的代码可以在 15 分钟内运行。

标签： pythonperformancefor-loopoptimizationnumba

解决方案

我找到的解决方案是使用 numba 内置的 parallel=True 装饰器。这将运行时间缩短到可接受的 20 分钟。这是最终结果的代码块。

 """
@njit(cache=True, nogil=True, fastmath=True, parallel=True)
def outerLoop(npref, fea, newFea, W):
    n = len(W)
    W4, W3, W2, W1 = W[n - 4], W[n - 3], W[n - 2], W[n - 1]
    for i in prange(len(fea)):

        aTerms = W4 + fea[i][0] * W3 + fea[i][1] * W2 + fea[i][2] * W1

        total = zeros((3,))

        for j in prange(len(npref)):
            dist = sum(square(npref[j] - fea[i]))
            u = dist * log(dist)
            total += u * W[j]

       
        newFea[i] = total + aTerms
    return newFea

推荐阅读