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

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

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

可是，现在老板胡远明说:

只需要导入一个叫Taichi库，可以提升100倍的代码速度！

不相信我。

看三个例子。

计算素数的个数，速度x120

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

标准答案如下:

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

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

现在，我们开始变魔术。

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

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

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

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

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

还不错。

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

动态规划，速度x500

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

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

例如，对于序列a =和序列b = ，它们的LCS是:

LCS = .

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

我们用f来表示这个子序列的长度，即LCS，前缀)。其中前缀(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。

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

其中有四个关键参数: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等等

当然，用太极和这些库和其他加速方式有什么区别胡远明也给出了详细的优缺点对比

支持0个人

反对

商品价值评分

快科技1997—2022版权所有。

ICP编号18024899 —2王编号41010502003949