k301
您当前的位置 :首页 > 科技 >

胡渊鸣:import一个“太极”库让Python代码提速100倍!

来源:网络    2022-09-13 20:38:00    阅读量:18448    
1_14 1_15 1_16
胡渊鸣:import一个“太极”库让Python代码提速100倍!

众所周知,Python的简单性和可读性是以牺牲性能为代价的mdashmdash

尤其是在计算密集型的情况下,例如多个for循环。

然而,现在老板胡远明说:

只需要导入一个叫ldquoTaichirdquo库,可以提升100倍的代码速度!

不相信我?

看三个例子。

计算素数的个数,速度x120

第一个例子非常非常简单。找出所有小于给定正整数n的质数。

标准答案如下:

我们将保存上面的代码并运行它。

当n为100万时,需要2.235秒才能得到结果:

现在,我们开始变魔术。

在不改变任何功能体的情况下,importldquotaichirdquo库,然后再添加两个装饰器:

答对了。同样的成绩只有0.363s,快了近6倍。

如果n = 1000万,只有0.8s;要知道,没有它也就55s,一下子快了70倍!

另外,我们可以给ti.init加一个参数,改成ti.init(arch=ti.gpu),这样taich就可以在gpu上进行计算了。

至此,计算所有小于1000万的质数只需要0.45s,比原来的Python代码快了120倍!

还不错?

什么?是不是觉得这个例子太简单,说服力不够?我们来看一个稍微复杂一点的。

动态规划,速度x500

不用说,动态规划作为一种优化算法,通过动态存储中间计算结果来减少计算时间。

我们使用经典教材《算法导论》中经典的动态编程案例ldquo最长公共子序列问题rdquo比如说。

例如,对于序列a =和序列b = (4,0,1,4,5,3,1,2),它们的LCS是:

LCS = (0,1,4,3,1,2).

用动态规划的思想计算LCS,就是先求出序列A的前I个元素和序列B的前J个元素的最长公共子序列的长度,然后逐渐增大I或J的值,重复这个过程得到结果。

我们用f来表示这个子序列的长度,即LCS((前缀(a,I),前缀(b,j))。其中前缀(A,I)表示序列A的前I个元素,即a(0),a(1),hellip,a(i-1),并且获得下面的递归关系:

完整的代码如下:

现在,我们用太极来加快速度:

结果如下:

胡远明电脑上的程序最快0.9秒就能完成,而用NumPy实现需要476秒,相差500多倍!

最后,我们来举一个不同的例子。

反应-扩散方程,效果惊人

在自然界中,总有一些动物的模式看似无序,其实并不完全随机。

艾伦·米德多特,图灵机的发明者;图灵是第一个提出模型来描述这种现象的人。

在该模型中,使用两种化学物质来模拟图案的产生。两者之间的关系类似于猎物和捕食者的关系,它们各自移动并相互作用:

最初,U和V随机分布在一个域中;

在每个时间步,它们逐渐扩散到相邻的空间;

当U和V相遇时,U的一部分被V吞噬,因此,V的浓度增加;

为了避免U被V根除,我们在每个时间步增加一定百分比的U,删除一定百分比(k)的V。

以上过程总结为ldquo反应扩散方程;:

其中有四个关键参数:Du、dv(v的扩散速率)、f(饲料的缩写,控制u的添加)和k(k的缩写,控制v的去除)。

如果在太极中实现这个方程,首先创建网格来表示域,用vec2来表示每个网格中U和V的浓度值。

拉普拉斯算子值的计算需要访问相邻网格。为了避免在同一个周期中更新和读取数据,我们应该创建两个形状相同的网格W x Htimes2。

每次从一个网格访问数据时,我们将更新的数据写入另一个网格,然后切换到下一个网格。然后数据结构设计是这样的:

首先,我们将网格中U的浓度设置为1,并将V放置在50个随机选择的位置:

那么实际的计算可以用不到10行代码来完成:

ti.kerneldefcompute:fori,jinti.ndrange(W,H):cen=uv(phase,I,j)lapl=uv(phase,i+1,j)+uv(phase,i-1,j)+uv(phase,I,J-1)-4.0 * cendu = du * lapl(0)-cen(0)* cen(1)* feed *(1-cen(0))DV = DV * lapl(1)+cen(0)* cen

最后一步是根据V的浓度对结果进行染色,可以得到这样一个惊人的图案:

有意思的是,胡远明介绍,即使V的初始浓度是随机设定的,每次也能得到类似的结果。

而相比Numba实现只能达到30fps左右,Taichi实现因为可以选择GPU作为后端,轻松超过300fps。

可以安装Pip install。

看了以上三个例子,你现在相信了吗?

实际上,Taichi是嵌入在Python中的DSL。它通过自己的编译器把ti.kernel修饰的函数编译到各种硬件上,包括CPU和GPU,然后进行高性能计算。

有了它,你再也不用羡慕C++/CUDA的性能了。

顾名思义,太极来自太极图形胡远明团队。现在只需使用pip install就可以安装这个库并与其他Python库交互,包括NumPy、Matplotlib和PyTorch等等。

当然,用太极和这些库和其他加速方式有什么区别?胡远明也给出了详细的优缺点对比。感兴趣的朋友可以戳下面的链接详细查看一下:

相关信息脑瘫男孩自学编程创业年收入300万网友:学会改变命运话到嘴边却想不起来该说什么:是痴呆的前兆吗?六七十岁的程序员工资暴涨!80/90无奈的原因:为什么会COBOL语言的人总想“接电话”?互联网巨头用什么编程语言进行开发?C++不可或缺,Java边缘化视图支持

支持0个人

反对

商品价值评分

快科技1997-2022版权所有。

ICP编号18024899 -2王编号41010502003949

声明:本网转发此文章,旨在为读者提供更多信息资讯,所涉内容不构成投资、消费建议。文章事实如有疑问,请与有关方核实,文章观点非本网观点,仅供读者参考。