import一个“太极”库,让python代码提速100倍!

Python小二 Python小二     2023-01-17     431

关键词:

丰色 发自 凹非寺
量子位 | 公众号 QbitAI

众所周知,Python的简单和易读性是靠牺牲性能为代价的——

尤其是在计算密集的情况下,比如多重for循环。

不过现在,大佬胡渊鸣说了:

只需import 一个叫做“Taichi”的库,就可以把代码速度提升100倍

不信?

来看三个例子。

计算素数的个数,速度x120

第一个例子非常非常简单,求所有小于给定正整数N的素数。

标准答案如下:

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

当N为100万时,需要2.235s得到结果:

现在,我们开始施魔法。

不用更改任何函数体,import“taichi”库,然后再加两个装饰器:

Bingo!同样的结果只要0.363s,快了将近6倍。

如果N=1000万,则只要0.8s;要知道,不加它可是55s,一下子又快了70倍

不止如此,我们还可以在ti.init()中加个参数变为ti.init(arch=ti.gpu) ,让taich在GPU上进行计算。

那么此时,计算所有小于1000万的素数就只耗时0.45s了,与原来的Python代码相比速度就提高了120倍

厉不厉害?

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

动态规划,速度x500

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

我们以经典教材《算法导论》中的经典动态规划案例“最长公共子序列问题(LCS)”为例。

比如对于序列a = [0, 1, 0, 2, 4, 3, 1, 2, 1]和序列b = [4, 0, 1, 4, 5, 3, 1, 2],它们的LCS就是:

LCS(a, b) = [0, 1, 4, 3, 1, 2]。

用动态规划的思路计算LCS,就是先求解序列a的前i个元素和序列b的前j个元素的最长公共子序列的长度,然后逐步增加i或j的值,重复过程,得到结果。

我们用f[i, j]来指代这个子序列的长度,即LCS((prefix(a, i), prefix(b, j)。其中prefix(a, i) 表示序列a的前i个元素,即a[0], a[1], …, a[i - 1],得到如下递归关系:

现在,我们用Taichi来加速:

结果如下:

胡渊鸣电脑上的程序最快做到了0.9秒内完成,而换成用NumPy来实现,则需要476秒,差异达到了超500倍!

最后,我们再来一个不一样的例子。

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

自然界中,总有一些动物身上长着一些看起来无序但实则并非完全随机的花纹。

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

在该模型中,两种化学物质(U和V)来模拟图案的生成。这两者之间的关系类似于猎物和捕食者,它们自行移动并有交互:

  1. 最初,U和V随机分布在一个域上;

  2. 在每个时间步,它们逐渐扩散到邻近空间;

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

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

上面这个过程被概述为“反应-扩散方程”:


其中有四个关键参数:Du(U的扩散速度),Dv(V的扩散速度),f(feed的缩写,控制U的加入)和k(kill的缩写,控制V的去除)。

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

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

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

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

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

@ti.kernel
def compute(phase: int):
    for i, j in ti.ndrange(W, H):
        cen = uv[phase, i, j]
        lapl = uv[phase, i + 1, j] + uv[phase, i, j + 1] + uv[phase, i - 1, j] + uv[phase, i, j - 1] - 4.0 * cen
        du = Du * lapl[0] - cen[0] * cen[1] * cen[1] + feed * (1 - cen[0])
        dv = Dv * lapl[1] + cen[0] * cen[1] * cen[1] - (feed + kill) * cen[1]
        val = cen + 0.5 * tm.vec2(du, dv)
        uv[1 - phase, i, j] = val

在这里,我们使用整数相位(0或1)来控制我们从哪个网格读取数据。

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

‍有趣的是,胡渊鸣介绍,即使V的初始浓度是随机设置的,但每次都可以得到相似的结果。

而且和只能达到30fps左右的Numba实现比起来,Taichi实现由于可以选择GPU作为后端,轻松超过了 300fps。

pip install即可安装

看完上面三个例子,你这下相信了吧?

其实,Taichi就是一个嵌入在Python中的DSL(动态脚本语言),它通过自己的编译器将被 @ti.kernel 装饰的函数编译到各种硬件上,包括CPU和GPU,然后进行高性能计算。

有了它,你无需再羡慕C++/CUDA的性能。

正如其名,Taichi就出自太极图形胡渊鸣的团队,现在你只需要用pip install就能安装这个库,并与其他Python库进行交互,包括NumPy、Matplotlib和PyTorch等等。

当然,Taichi用起来和这些库以及其他加速方法有什么差别,胡渊鸣也给出了详细的优缺点对比,感兴趣的朋友可以戳下面的链接详细查看:

https://docs.taichi-lang.org/blog/accelerate-python-code-100x

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

...多重for循环。不过现在,大佬胡渊鸣说了:只需import一个叫做“Taichi”的库,就可以把代码速度提升100倍!不信?来看三个例子。计算素数的个数,速度 查看详情

几个压箱底的tips,python提速不是一点点(代码片段)

...优化细节之前,需要了解一些代码优化基本原则。第一个基本原则是不要过早优化。很多人一开始写代码就奔着性能优化的目标,“让正确的程序更快要比让快速的程序正确容易得多”。因此&#x 查看详情

几个压箱底的tips,python提速不是一点点(代码片段)

...优化细节之前,需要了解一些代码优化基本原则。第一个基本原则是不要过早优化。很多人一开始写代码就奔着性能优化的目标,“让正确的程序更快要比让快速的程序正确容易得多”。因此&#x 查看详情

import详解(代码片段)

...hon解释器能够搜索到这个模块,这种方式比较灵活。查看import搜索路径importsysforiinsys.path:print(i)设置搜索路径im 查看详情

python调用c语言实现数独计算逻辑提速100倍(代码片段)

...icle/details/118583461最终达到了这样的效果:中间我们将一个超难的数独,程序计算耗时由17秒优化到3秒,有读者希望将计算数独 查看详情

python破解wifi密码怎么提速?

...能的密码,有时候你可能知道几个密码,但是不确定是哪一个了,你可以把预期的密码,写在这里面,让程序自己去试验。我个人最常用这个功能,尤其是在切换不同wifi的时候如果你完全不知道密码有多少位,可能的值是什么,... 查看详情

python中标准模块importlib详解(代码片段)

Python中标准模块importlib详解模块简介Python提供了importlib包作为标准库的一部分。目的就是提供Python中import语句的实现(以及__import__函数)。另外,importlib允许程序员创建他们自定义的对象,可用于引入过程(也称为importer)。什... 查看详情

用taichi加速python:提速100+倍!(代码片段)

Python已经成为世界上最流行的编程语言,尤其在深度学习、数据科学等领域占据主导地位。但是由于其解释执行的属性,Python较低的性能很影响它在计算密集(比如多重for循环)的场景下发挥作用,实在让人... 查看详情

优化nginxhttps延迟-如何让nginx提速30%的?(代码片段)

...alancer)、反向代理(ReverseProxy),以及网关(Gateway)等等。一个配置得当的Nginx 查看详情

一个让gitclone提速几十倍的小技巧(实用!)

...会失败的那种。大家会怎么处理的呢?可能会考虑换一个下载源,可能会通过一些手段提高网速,但是如果这些都试过了还是比较慢呢?今天我就遇到了这个问题,我需要把typescript代码从gitlab下载下来,... 查看详情

你敢信!?几行代码让swift数组初始化提速440+倍!

...需求在Swift中,我们有时需要初始化数组内容,这往往是一个耗时操作,数组越大时间越长。那么,如何最大限度的缩短数组的初始化时间呢?如上所示,在初始化1w个元素的数组时,我们的3种方法分别耗时为:53.855秒0.36秒0.12... 查看详情

你敢信!?几行代码让swift数组初始化提速400多倍!

...需求在Swift中,我们有时需要初始化数组内容,这往往是一个耗时操作,数组越大时间越长。那么,如何最大限度的缩短数组的初始化时间呢?如上所示,在初始化1w个元素的数组时,我们的3种方法分别耗时为:53.855秒0.36秒0.12... 查看详情

import的使用(代码片段)

首先来看一个例子:importmathprint(math.pi)#输出3.141592653589793上面的程序使用了import语句,importmath的意思是从Python标准库中引入math.py模块,这是Python中定义的引入模块的方法可以用如下方法引入多个模块:importmath,time,sys除了用import... 查看详情

如何通过执行sql为低代码项目提速?

...按照人类的需求进行工作,满足人类的需要就成为了一个问题,于是便诞生了计算机语言。最初的计算机语言是由0和1构成的机器语言& 查看详情

如何通过执行sql为低代码项目提速?

...按照人类的需求进行工作,满足人类的需要就成为了一个问题,于是便诞生了计算机语言。最初的计算机语言是由0和1构成的机器语言& 查看详情

如何通过执行sql为低代码项目提速?

...按照人类的需求进行工作,满足人类的需要就成为了一个问题,于是便诞生了计算机语言。最初的计算机语言是由0和1构成的机器语言& 查看详情

python二级考试知识点(史上最全)(代码片段)

...值语句、库引用33个保留字6种数据类型4种引用方法:import库、from库import函数、from库impor*、import库as别名考点1.2基本输入输出函数:input 查看详情

dart库(代码片段)

...分片partTOCDart库Dart中的库主要有三种:1、我们自定义的库import‘lib/xxx.dart;import‘../../part-two/chapter3/enum.dart‘;2、系统内置库impont‘dart:mathimpont‘dant:io‘;import‘dart:convert‘;3、Pub包管理系统中的库(https://pub.dev)https://pub.dev/packageshttp... 查看详情