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

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

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

在我的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如果你想要精确(或至少不是太不精确)的结果,我已经说过可能不合适。

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

热门问答

在serverless中,我能否自己host 一个express(nodejs)的服务?

Tina

腾讯云 · 产品经理 (已认证)

Go Serverless!
推荐
您好,可以这样的。您可以参考如下文档,申请下http function 您可以使用常见的 WEB 框架(如 Nodejs Web 框架:Express、Koa)编写 HTTP 函数。而 WEB 框架内置的一些中间件(如cors)也会极大的方便您的业务编写 文档链接 https:...... 展开详请

腾讯云安全方面,获得过哪些认证?

腾讯安全

腾讯 · 腾讯安全 (已认证)

腾讯安全,致力于成为产业数字化升级的安全战略官,守护政企信息安全,为产业数字化升级保驾护航。
推荐
你好,合规性是腾讯云发展的基础,腾讯安全助力腾讯云,满足不同行业、领域、国家的合规性要求,全力打造值得客户信赖的云服务;同时,积极参与行业安全标准的制定及推广,坚持合规即服务,建设和运行安全可靠的云生态环境。 腾讯云安全目前获得的认证,包括但不限于以下21项。 国际权威 腾讯...... 展开详请

使用有过期时间的签名往Cos存储桶中上传文件,若上传还在进行中签名过期,上传是否会终止?

galenye

腾讯 · 工程师 (已认证)

对象存储专业搬砖工
推荐已采纳

如果你是使用的简单上传,它能接收5g以内的文件,那签名过期的文件还在上传的话,是没影响的,因为签名判断是在cos接受到请求时。

如果你是使用的sdk等封装的分片上传,那其实是多个请求去上传文件,如果签名过期了,那上传到某一刻,后面的请求都会返回403

存储桶的默认加速域名 cdn 如何更改业务类型, 即把静态加速改成下载加速?

Jinqn

腾讯 · 高级工程师 (已认证)

腾讯云COS前端开发
推荐

我理解你意思是,浏览器打开的时候要下载,不要直接显示。

通过存储桶的文件 Content-Type 来控制

为何我使用.Net API 生成的临时密钥无法进行文件操作?

推荐
cos有自己的密钥系统,应该是在控制台上,访问管理,API密钥,项目密钥那里,或者去看看cos的文档是如何说明的吧。 你通过ms接口创建cos临时密钥,也许的确会被限制一些,这个需要ms这个产品的人回答下比较好。 生成临时密钥和哪个SDK无关,可以直接在线调用也可以生成,通过AP...... 展开详请

tencentcloud-sdk-php-master github代码上没有vendor文件夹?

推荐
因为和composer冲突,因此导出时没有包含vendor目录。如果需要,可以考虑git clone方式拿到,或者到https://cloud.tencent.com/document/sdk/PHP#.E9.80.9A.E8.BF.87.E6.BA.90.E7.A0.81.E5...... 展开详请

所属标签

扫码关注云+社区

领取腾讯云代金券