可以优化numba函数中的循环以更快地运行吗?

内容来源于 Stack Overflow,并遵循CC BY-SA 3.0许可协议进行翻译与使用

  • 回答 (2)
  • 关注 (0)
  • 查看 (50)

在我的python代码中,我需要循环大约2500万次,我想尽可能地优化。循环中的操作非常简单。为了使代码高效,我使用了numba模块,这有很大帮助,但如果可能的话,我想进一步优化代码。

这是一个完整的工作示例:

import numba as nb
import numpy as np
import time 
#######create some synthetic data for illustration purpose##################
size=5000
eps = 0.2
theta_c = 0.4
temp = np.ones(size)
neighbour = np.random.randint(size, size=(size, 3)) 
coschi = np.random.random_sample((size))
theta = np.random.random_sample((size))*np.pi/2
pwr = np.cos(theta)
###################end of dummy data##########################

###################-----main loop------###############
@nb.jit(fastmath=True)
def func(theta, pwr, neighbour, coschi, temp):
    for k in range(np.argmax(pwr), 5000*(pwr.size)): 
        n = k%pwr.size
        if (np.abs(theta[n]-np.pi/2.)<np.abs(theta_c)):
                adj = neighbour[n,1]
        else:
                adj = neighbour[n,0]
        psi_diff = np.abs(np.arccos(coschi[adj])-np.arccos(coschi[n]))
        temp5 = temp[adj]**5;
        e_temp = 1.- np.exp(-temp5*psi_diff/np.abs(eps))
        temp[n] = temp[adj] + (e_temp)/temp5*(pwr[n] - temp[adj]**4)
    return temp

#check time
time1 = time.time()
temp = func(theta, pwr, neighbour, coschi, temp)
print("Took: ", time.time()-time1, " seconds.")

3.49 seconds取决于我的机器。

我需要运行这个代码几千次来进行一些模型拟合,因此优化甚至1秒意味着为我节省了数十个小时。

如何进一步优化此代码?

提问于
用户回答回答于

让我先谈谈一些一般性意见:

  • 如果您使用numba并且非常关心性能,则应该避免numba创建对象模式代码的任何可能性。这意味着你应该使用numba.njit(...)numba.jit(nopython=True, ...)代替numba.jit(...)。 这对你的情况没有任何影响,但它会使意图更清晰,并且在(快速)nopython模式中不支持某些内容时会抛出异常。
  • 你应该小心你的时间和方式。对numba-jitted函数的第一次调用(未提前编译)将包括编译成本。因此,您需要在计时之前运行一次以获得准确的计时。要获得更准确的计时,您应该多次调用该函数。我喜欢Jupyter %timeit笔记本/实验室中的IPythons来了解性能。 所以我会用: res1 = func(theta, pwr, neighbour, coschi, np.ones(size)) res2 = # other approach np.testing.assert_allclose(res1, res2) %timeit func(theta, pwr, neighbour, coschi, np.ones(size)) %timeit # other approach 这样我就使用第一次调用(包括编译时间)和一个断言,以确保它确实产生(几乎)相同的输出,然后使用更强大的定时方法(与之相比time)对函数计时。

Hoisting np.arccos

现在让我们从一些实际的性能优化开始:一个显而易见的是你可以提升一些“不变量”,例如,np.arccos(coschi[...])计算的频率比实际的元素要多得多coschi。你在coschi大约5000次迭代每个元素,它np.arccos每循环做两次!因此,让我们计算arccoscoschi一次,并将其存储在一个中间数组这样一个可以访问内循环:

@nb.njit(fastmath=True)
def func2(theta, pwr, neighbour, coschi, temp):
    arccos_coschi = np.arccos(coschi)
    for k in range(np.argmax(pwr), 5000 * pwr.size): 
        n = k % pwr.size
        if np.abs(theta[n] - np.pi / 2.) < np.abs(theta_c):
            adj = neighbour[n, 1]
        else:
            adj = neighbour[n, 0]
        psi_diff = np.abs(arccos_coschi[adj] - arccos_coschi[n])
        temp5 = temp[adj]**5;
        e_temp = 1. - np.exp(-temp5 * psi_diff / np.abs(eps))
        temp[n] = temp[adj] + e_temp / temp5 * (pwr[n] - temp[adj]**4)
    return temp

在我的计算机上已经快得多:

1.73 s ± 54.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # original
811 ms ± 49.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # func2

然而它需要付出代价:结果会有所不同!我使用原始版本和提升版本始终获得显着不同的结果fastmath=True。然而,结果(几乎)相等fastmath=False。这似乎可以fastmath实现一些严格的优化,np.arccos(coschi[adj]) - np.arccos(coschi[n])np.arccos提升时是不可能的。在我个人看来,fastmath=True如果您关心确切的结果,或者您已经测试过结果的准确性不会受到fastmath的显着影响,我会忽略它!

Hoisting adj

提升的下一件事就是adj,它的计算频率也超过了必要的程度:

@nb.njit(fastmath=True)
def func3(theta, pwr, neighbour, coschi, temp):
    arccos_coschi = np.arccos(coschi)
    associated_neighbour = np.empty(neighbour.shape[0], nb.int64)
    for idx in range(neighbour.shape[0]):
        if np.abs(theta[idx] - np.pi / 2.) < np.abs(theta_c):
            associated_neighbour[idx] = neighbour[idx, 1]
        else:
            associated_neighbour[idx] = neighbour[idx, 0]

    for k in range(np.argmax(pwr), 5000 * pwr.size): 
        n = k % pwr.size
        adj = associated_neighbour[n]
        psi_diff = np.abs(arccos_coschi[adj] - arccos_coschi[n])
        temp5 = temp[adj]**5;
        e_temp = 1. - np.exp(-temp5 * psi_diff / np.abs(eps))
        temp[n] = temp[adj] + e_temp / temp5 * (pwr[n] - temp[adj]**4)
    return temp

这种影响并不大,但可以衡量:

1.75 s ± 110 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # original
761 ms ± 28.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func2
660 ms ± 8.42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func3

提升额外的计算似乎对我的计算机上的性能没有影响,因此我不在此处包含它们。所以看来,如果不改变算法,你可以获得多远。

重构为更小的函数(+小的额外更改)

但是我建议将其他函数中的提升分开并使所有变量都是函数参数而不是查找全局变量。这可能不会导致加速,但它可以使代码更具可读性:

@nb.njit
def func4_inner(indices, pwr, associated_neighbour, arccos_coschi, temp, abs_eps):
    for n in indices:
        adj = associated_neighbour[n]
        psi_diff = np.abs(arccos_coschi[adj] - arccos_coschi[n])
        temp5 = temp[adj]**5;
        e_temp = 1. - np.exp(-temp5 * psi_diff / abs_eps)
        temp[n] = temp[adj] + e_temp / temp5 * (pwr[n] - temp[adj]**4)
    return temp

@nb.njit
def get_relevant_neighbor(neighbour, abs_theta_minus_pi_half, abs_theta_c):
    associated_neighbour = np.empty(neighbour.shape[0], nb.int64)
    for idx in range(neighbour.shape[0]):
        if abs_theta_minus_pi_half[idx] < abs_theta_c:
            associated_neighbour[idx] = neighbour[idx, 1]
        else:
            associated_neighbour[idx] = neighbour[idx, 0]
    return associated_neighbour

def func4(theta, pwr, neighbour, coschi, temp, theta_c, eps):
    arccos_coschi = np.arccos(coschi)
    abs_theta_minus_pi_half = np.abs(theta - (np.pi / 2.))
    relevant_neighbors = get_relevant_neighbor(neighbour, abs_theta_minus_pi_half, abs(theta_c))
    argmax_pwr = np.argmax(pwr)
    indices = np.tile(np.arange(pwr.size), 5000)[argmax_pwr:]
    return func4_inner(indices, pwr, relevant_neighbors, arccos_coschi, temp, abs(eps))

在这里,我还做了一些额外的改动:

  • 预先使用np.tile和切片计算指数而不是range方法%
  • 使用普通NumPy(在numba之外)来计算np.arccos

最后时间和总结

1.79 s ± 49.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # original
844 ms ± 41.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # func2
707 ms ± 31.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # func3
550 ms ± 4.88 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # func4

因此,最终的方法fastmath比原始方法快大约3倍(没有)。如果你确定要使用的fastmath,那么就适用fastmath=Truefunc4_inner,这将是更快:

499 ms ± 4.47 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # func4 with fastmath on func4_inner 

但是,fastmath如果你想要精确(或至少不是太不精确)的结果,我已经说过可能不合适。

此外,这里的一些优化很大程度上取决于可用的硬件和处理器缓存(特别是对于代码的内存带宽限制部分)。您必须检查这些方法在计算机上相对于彼此的执行情况。

用户回答回答于

对于初学者,您可以使用标准库和数学函数而不是numpy来执行某些操作。仅这些变化在我的计算机上从2.2秒到1.9秒。

import numba as nb
import numpy as np
import math
import time 
#######create some synthetic data for illustration purpose##################
size=5000
eps = 0.2
theta_c = 0.4
temp = np.ones(size)
neighbour = np.random.randint(size, size=(size, 3)) 
coschi = np.random.random_sample((size))
theta = np.random.random_sample((size))*np.pi/2
pwr = np.cos(theta)
###################end of dummy data##########################

###################-----main loop------###############
@nb.jit(fastmath=True)
def func(theta, pwr, neighbour, coschi, temp):
    for k in range(np.argmax(pwr), 5000*(pwr.size)): 
        n = k%pwr.size
        #taking into account regions with different super wind direction
        if (abs(theta[n]-math.pi/2.)<abs(theta_c)):
                adj = neighbour[n,1]
        else:
                adj = neighbour[n,0]
        psi_diff = abs(math.acos(coschi[adj])-math.acos(coschi[n]))
        temp5 = temp[adj]**5;
        e_temp = 1.- math.exp(-temp5*psi_diff/abs(eps))
        temp[n] = temp[adj] + (e_temp)/temp5*(pwr[n] - temp[adj]**4)
    return temp

#check time
time1 = time.time()
temp = func(theta, pwr, neighbour, coschi, temp)
print("Took: ", time.time()-time1, " seconds.")

扫码关注云+社区

领取腾讯云代金券