我正在编写将变换应用于一组 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 分钟内运行。
答案 0 :(得分:0)
我找到的解决方案是使用 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