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

Python猫 Python猫     2022-11-30     438

关键词:

Python 已经成为世界上最流行的编程语言,尤其在深度学习、数据科学等领域占据主导地位。但是由于其解释执行的属性,Python 较低的性能很影响它在计算密集(比如多重 for 循环)的场景下发挥作用,实在让人又爱又恨。如果你是一名经常需要使用 Python 进行密集计算的开发者,我相信你肯定会有下面的类似经历:

  • 我的 Python 程序里面有个很大的 for 循环,循环体里面全是密集的计算,跑起来好慢啊...

  • 我的程序里面只有一小部分计算是性能瓶颈,虽然可以用 C++ 改写然后用 ctypes 绑定一下,但是那样会很麻烦,还会有在别的机器上编译不了的风险。我希望所有的工作都能在一个 Python 脚本中完成!

  • 我之前是忠实的 C++/Fortran 用户,但是最近周围的同学用 Python 的越来越多,我也想试试 Python,但是无奈很多祖传代码用 Python 改写以后就会慢 100 多倍,我接受不了...

  • 我的工作中需要处理大量图片数据,而需要的图像处理功能 OpenCV 又不提供,只能自己手写两重 for 循环,在 Python 里面这么搞真是太痛苦了 ...

如果你有类似的烦恼,那真的值得了解一下 Taichi。我来简单介绍一下:Taichi 是一个嵌入在 Python 中的领域特定语言,其一大功能就是加速 Python,让 Python 代码跑得和 C++ 甚至 CUDA 一样快。Taichi 通过自己的编译器将被 @ti.kernel 修饰的函数编译到各种硬件上,包括 CPU 和 GPU,然后高性能执行。

(用户不用关心的)Taichi 运行原理:Python 代码被 Taichi 编译器编译到高性能二进制

由于 Taichi 开发者社区花了大量的精力优化 Taichi 在 Python 中的使用体验,所有的 Taichi 功能都可以在 import taichi as ti 以后使用,Taichi 本身也可以使用 pip 进行安装。当然,Taichi 也可以与常用的 Python 包(numpy、matplotlib、PyTorch 等)进行交互。

在这篇文章中,我们将通过三个计算例子来演示如何使用 Taichi 让你的 Python 轻松加速 > 50 倍。这三个例子是:1. 计算质数数目;2. 动态规划求解最长公共子序列;3. 求解反应-扩散方程。

🔗 https://github.com/taichi-dev/faster-python-with-taichi

计算素数个数

作为开胃小菜,我们先做一个小实验:计算小于给定正整数 的素数的个数。相信任何对 Python 有基础了解的人都不难写出类似下面这样的解法:

"""Count the number of primes in range [1, n].
"""

def is_prime(n: int):
    result = True
    for k in range(2, int(n ** 0.5) + 1):
        if n % k == 0:
            result = False
            break
    return result

def count_primes(n: int) -> int:
    count = 0
    for k in range(2, n):
        if is_prime(k):
            count += 1

    return count

print(count_primes(1000000))

这个方法的思路简单且粗暴:我们用一个函数 is_prime 来判断某个正整数 是不是素数,是素数则返回 1,不是则返回 0。这只要遍历检查从 2 到 之间是否有整数能够整除 即可。然后将小于 的全部整数依次代入此函数并统计结果。将上面的代码保存为 count_primes.py,在命令行运行:

time python count_primes.py

在我的电脑上输出的运行结果是:

78498

real        0m2.235s
user        0m2.235s
sys        0m0.000s

耗时 2.235 秒。也许代码中 设置成一百万对你的电脑来说太轻松了,要不要把 改成一千万试试?我打赌不管你的电脑多么高端,你起码都要等个半分钟才能看到结果。

好了下面是魔法时刻:我们不修改上面的函数体,只 import 一个“库”,然后给两个函数分别加一个装饰器:

"""Count the number of primes below a given bound.
"""
import taichi as ti
ti.init()

@ti.func
def is_prime(n: int):
    result = True
    for k in range(2, int(n ** 0.5) + 1):
        if n % k == 0:
            result = False
            break
    return result

@ti.kernel
def count_primes(n: int) -> int:
    count = 0
    for k in range(2, n):
        if is_prime(k):
            count += 1

    return count

print(count_primes(1000000))

仍然运行 time python count_primes.py 命令,输出的结果是:

78498

real        0m0.363s
user        0m0.546s
sys        0m0.179s

速度直接 x6! 而将 改成一千万的话,Taichi 的耗时只会增加到 0.8s 左右,而 Python 则需要大约 55 秒,Taichi 直接加速了 70 倍!不仅如此,我们还可以在 ti.init 中加上 ti.init(arch=ti.gpu) 参数,指定 Taichi 使用 GPU 来进行计算。在 GPU 上同样的计算 Taichi 只花了不到 0.45 秒,比 Python 足足快了 120 倍!你可以运行这里的代码亲身体会一下。

上面这个计算素数的例子使用的方法有点土,作为习题还可以,但在实际生产中就显得不那么实用了。我们接下来看一个实际中普遍使用的算法。

动态规划

动态规划(Dynamic Programming)是一类特别实用的算法,这类算法的哲学是以空间换时间,通过存储中间计算结果来减少重复计算量。我们这里选择一个求解最长公共子序列(Longest common subsequence, LCS)的例子 (算法导论的读者有木有)。

插播两个来自渊鸣的《算法导论》小故事:

  1. 笔者小时候买过一本《算法导论》,书中提到四位作者都来自“麻雀理工学院”。当时还很好奇:怎么会有学校叫这么奇怪的名字... 过了一阵才意识到自己可能成了盗版书籍的受害者。

  2. 10 年后,我还真的来到了麻省理工学院(MIT)读博士,一年后进行硕士论文答辩(MIT 叫做 Research Qualification Exam),我自然就带着 Taichi 的论文去了。答辩委员会里面有一位慈祥的教授,Charles E. Leiserson,嗯,就是《算法导论》的作者之一,“CLRS” 之中的 L。

言归正传。所谓子序列,就是一个序列的子集,但是保持它们在原序列中的顺序。比如说 [1, 2, 1][1, 2, 3, 1] 的子序列,而 [3, 2] 则不是。我们这里考虑对两条给定的序列,求出它们最长公共子序列的长度。最长公共子序列就是两个序列的所有公共子序列中最长的一条 (这个最长子序列未必唯一,但它的长度是唯一确定的)。

举个例子:

a = [0, 1, 0, 2, 4, 3, 1, 2, 1]

b = [4, 0, 1, 4, 5, 3, 1, 2]

的最长公共子序列是

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

最长公共子序列有很多应用。比如大家日常使用的 Linux diff 命令和 git 工具(比较两个文件之间的相似度),还有生物信息学中判断两段基因的相似度(把数字换成 ACGT 就行),其中的实现都用到了 LCS。

动态规划计算 LCS 的想法是我们依次求解序列 a 的前 i 个元素和序列 b 的前 j 个元素的最长公共子序列的长度,通过让 ij 逐渐增加我们就逐步得出了最终的结果。我们用 f[i, j] 表示 LCS((prefix(a, i), prefix(b, j),其中 prefix(a, i) 表示序列 a 的前 i 个元素,即 a[0], a[1], ..., a[i - 1]。这样我们就得到递推式:

f[i, j] = max(f[i - 1, j - 1] + (a[i - 1] == b[j - 1]),
              max(f[i - 1, j], f[i, j - 1]))

于是,一个 LCS 算法用 Python 可以很自然地书写为:

for i in range(1, len_a + 1):
    for j in range(1, len_b + 1):
        f[i, j] = max(f[i - 1, j - 1] + (a[i - 1] == b[j - 1]),
                      max(f[i - 1, j], f[i, j - 1]))

这里我们给出一个 Taichi 的加速实现:

import taichi as ti
import numpy as np

ti.init(arch=ti.cpu)

benchmark = True

N = 15000

f = ti.field(dtype=ti.i32, shape=(N + 1, N + 1))

if benchmark:
    a_numpy = np.random.randint(0, 100, N, dtype=np.int32)
    b_numpy = np.random.randint(0, 100, N, dtype=np.int32)
else:
    a_numpy = np.array([0, 1, 0, 2, 4, 3, 1, 2, 1], dtype=np.int32)
    b_numpy = np.array([4, 0, 1, 4, 5, 3, 1, 2], dtype=np.int32)


@ti.kernel
def compute_lcs(a: ti.types.ndarray(), b: ti.types.ndarray()) -> ti.i32:
    len_a, len_b = a.shape[0], b.shape[0]

    ti.loop_config(serialize=True) # 避免 Taichi 自动并行
    for i in range(1, len_a + 1):
        for j in range(1, len_b + 1):
            f[i, j] = max(f[i - 1, j - 1] + (a[i - 1] == b[j - 1]),
                          max(f[i - 1, j], f[i, j - 1]))

    return f[len_a, len_b]


print(compute_lcs(a_numpy, b_numpy))

将上面的代码保存为 lcs.py,然后在终端运行:

time python lcs.py

得到的结果为(具体结果每次未必一致):

2721

real        0m1.409s
user        0m1.112s
sys        0m0.549s

我们在代码中同时提供了分别使用 Taichi 和 Numpy 计算的版本,在我的电脑上对两个长度是 N=15000 的随机序列进行计算 Taichi 版本大约需要 0.9 秒,而 Python 则需要 476s,足足差了 500 多倍!大家可以运行一下体会 Taichi 相对 Numpy 那种飞一样的感觉。

当然,Numpy 主要针对的场景是以数组为基本单位的运算,遇到这种需要在数组内更细粒度进行计算的情况就比较无力了。而这正是 Taichi 能够发挥作用的地方。

反应 - 扩散方程

在大自然中我们常常会在动植物的表面见到一些有趣的图案,比如斑马身上的条纹,猎豹身上的斑点,河豚表面的花纹等等。

这些图案看起来是不规则的,但是又有一定的规律,并不完全随机。从进化的观点,这些图案是生物在长期演进和自然选择中逐渐形成的,但到底是什么规则决定了它们的形状一直是个有趣的问题。阿兰 . 图灵 (正是图灵机的发明人) 是最早注意到这一现象并尝试给出模型描述的人。他在论文 "The Chemical Basis of Morphogenesis" 中提出可以用两种化学物质 U, V 之间的相互作用来模拟图案的形成过程,其中物质 U 的角色类似被捕食者 (prey),物质 V 的角色类似捕食者 (predator)。它们之间的作用服从如下规则:

  1. 初始时空间中随机地分布了一些 U, V。

  2. 在每个时刻 1, 2, 3, ..., U, V 两种物质都向其邻域扩散。

  3. 当 U, V 相遇时,一定比例的 U, V 会合并转化为更多的 V (捕食者在捕食后数量会增加)

  4. 为了避免捕食者 V 的数量过多导致 U 的数量被消耗光,我们在每个时刻按照一定的比例 f 添加 U,同时按照一定的比例 k 移走 V。

于是整个过程可以用下面的反应 - 扩散方程描述:

这里关键的控制参数有四个,分别是 , 分别控制 U, V 的扩散速度, 代表 feed,控制 U 的添加量,而 代表 kill,控制移走 V 的比例。

为了在 Taichi 中模拟这一过程,我们将空间划分为网格,每个网格中 U, V 的浓度值用一个 vec2 来表示。注意拉普拉斯算子 的数值计算是需要访问当前网格周围的网格的,为了避免一边修改一边读取这种操作的发生,我们需要开辟两个形状为 的网格,每次用其中一个网格的值作为旧值,将更新后的浓度值写入另一个网格中,然后交换两个网格的角色。所以我们需要的数据结构应该是:

W, H = 800, 600
uv = ti.Vector.field(2, float, shape=(2, W, H))

初始时,我们假定网格中 U 的浓度处处是 1,然后随机选择 50 个点撒上 V:

import numpy as np

uv_grid = np.zeros((2, W, H, 2), dtype=np.float32)
uv_grid[0, :, :, 0] = 1.0
rand_rows = np.random.choice(range(W), 50)
rand_cols = np.random.choice(range(H), 50)
uv_grid[0, rand_rows, rand_cols, 1] = 1.0
uv.from_numpy(uv_grid)

实际的计算代码非常之简短:

@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 的整数 phase 来控制使用 uv 的哪一层来作为旧的网格,并将更新的值写入 1-phase 对应的层中。

根据 V 的浓度进行染色,我们得到了如下的动画效果:

非常有趣的是,虽然 V 的初始浓度是随机设置的,但是最终得到的图案却具有相似性。

我们在代码中提供了基于 Taichi 和 Numba 的两份不同的实现,Taichi 的版本由于使用了 GPU 进行计算,计算的部分可以轻松达到 300+ fps,而 Numba 的版本计算部分虽然也是编译执行的,但由于是在 CPU 上计算的,只有大约 30fps 左右。大家可以亲自运行代码体会一下 Taichi 使用 GPU 加速的巨大优势。

总结

在这三个例子上 Taichi 都让程序有了大幅加速。主要的性能来自三点:

  1. Taichi 是编译性的,而 Python 是解释性的

  2. Taichi 能自动并行,而 Python 通常是单线程的

  3. Taichi 能在 GPU 上运行,而 Python 本身是在 CPU 上运行的

当然,加速 Python 还有很多其他工具,这里我们分析一下他们和 Taichi 的优劣。

与 Numpy/JAX/PyTorch/TensorFlow 比较:这几类工具都高度基于数组运算。计算的最小单位是数组,在 Data Science、Deep Learning 等领域是有明显的优势的。但是在科学计算领域,这样做导致灵活性缺失:比如说前面那个计算质数的程序,就比较难使用数组运算表示出来。Taichi 的优势就在于其灵活性,能够直接操纵循环的每一次迭代,以一种更细颗粒度进行对于计算的描述,类似 C++ 和 CUDA。

与 Cython 比较:使用 Cython 编写程序实现加速也是一种常见的选择。在 Numpy 和 Scipy 的官方代码中有不少模块都是使用 Cython 编写然后编译的。但按照 Cython 的要求书写代码会比较麻烦,会牺牲一些可读性。Cython 支持一定程度的并行计算,但不支持直接调用 GPU 进行计算。

与 Numba 比较:Numba 顾名思义,是非常适合针对 Numpy 进行加速的方案。当你的函数是针对 Numpy 的数组向量化的操作时,使用 Numba 将其编译以后执行可以大大加速。Taichi 相比 Numba 的优势还有:1. Taichi 支持各种灵活的数据类型,比如 struct, dataclass, quant, sparse 等等,你可以任意指定它们的内存排布,当数据量庞大时这个优势会非常明显。而 Numba 只有在针对 Numpy 的稠密数组时效果最佳。2. Taichi 可以调用不同的 GPU 后端进行计算,所以写大规模并行程序(如粒子仿真、渲染器等)这种操作对 Taichi 来说是小菜一碟。但你很难想象可以用 Numba 写一个还过得去的 (哪怕离线) 渲染器。

与 Pypy 比较:Pypy 是一个 Python 的 JIT 编译器,这个工具 2007 年就有了,和 Taichi 的解决方案有些类似,都是通过编译的方式加速 Python。Pypy 最大优势在于 Python 代码完全不用改变,就能通过 Pypy 加速。但是这也是 Pypy 加速比率比 Taichi 低的原因:因为 Pypy 需要在编译的同时保持 Python 所有的语言特性,所以能够进行的优化比较有限。而 Taichi 有一套自己的语法,虽然和 Python 很像但是也有自己的一些假设,这使得 Taichi 能够实现更大的加速。

与 ctypes 比较:ctypes 可以让用户在 Python 中调用 C 函数。C++、CUDA 编写的程序也可以用过 C 接口暴露给 Python 使用。但是,ctypes 会让工具链复杂化:为了写一段比较快的程序,用户需要同时掌握 C、Python、CMake、CUDA 等等语言,和本文描述的完全在 Python 中解决问题的方案比起来还是麻烦了一些。

总而言之,在科学计算任务上,Taichi 还是有自己独特的优势的,大家可以根据自己的需求选择最合适的工具。如果你需要在 Python 中实现类似 C/C++ 语言的性能,Taichi 不失为一个理想的选择!

最后,我们希望 Taichi 能够为你带来价值,也希望能够听到你对 Taichi 的反馈,欢迎给我们提交 issues,加入 Taichi 开发者社区。如果想一键体验 Taichi,只需要执行👇🏻

pip install -U taichi

并执行👇🏻

ti gallery

就可以体验各种基于 Taichi 的高性能可视化 Demos,期待与大家相遇!

Python猫技术交流群开放啦!群里既有国内一二线大厂在职员工,也有国内外高校在读学生,既有十多年码龄的编程老鸟,也有中小学刚刚入门的新人,学习氛围良好!想入群的同学,请在公号内回复『交流群』,获取猫哥的微信(谢绝广告党,非诚勿扰!)~

还不过瘾?试试它们

7 行代码搞崩溃 B 站,原因令人唏嘘!

收钱吧的 Python 高效自动化测试实践

重写 500 Lines or Less 项目:流水线调度问题

把Redis当作队列来用,真的合适吗?

14K 字,25 张图彻底掌握分布式事务原理!

Python 的 heapq 模块源码分析

如果你觉得本文有帮助

请慷慨分享点赞,感谢啦

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

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

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

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

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

之前我们实现python自动玩数独游戏,详见:《让程序自动玩数独游戏让你秒变骨灰级数独玩家》链接:https://blog.csdn.net/as604049322/article/details/118583461最终达到了这样的效果:中间我们将一个超难的数独,程序计... 查看详情

如何让xcode在读写上提速100倍?

如何让Xcode在读写上提速100倍?上个月参加了一场西雅图当地的线下iOS开发者聚会。JeffSzuhay作为一个有20+年开发经验的资深程序员,跟我讲了一套提高iOS开发效率的方法。相比于其他程序员在App启动时间、架构优化方面的经验,... 查看详情

深度学习提速10倍!

点上方人工智能算法与Python大数据获取更多干货在右上方 ··· 设为星标 ★,第一时间获取资源仅做学术分享,如有侵权,联系删除转载于:量子位面对数以亿计的图片数据,到底该用什么样的方法才能快... 查看详情

深度学习提速10倍!

点上方人工智能算法与Python大数据获取更多干货在右上方 ··· 设为星标 ★,第一时间获取资源仅做学术分享,如有侵权,联系删除转载于:量子位面对数以亿计的图片数据,到底该用什么样的方法才能快... 查看详情

python之父爆料:明年至少令python提速1倍!(代码片段)

...程序员交流群????????作者:豌豆花下猫  来源:Python猫  大概在半年前,我偶然看到一篇文章,有人提出了给Python提速5倍的计划,并在寻找经费赞助。当时并没有在意,此后也没有看到这方面的消息。但... 查看详情

快出数量级的性能是怎样炼成的

...能优化的案例,不算很多,还没有失手过。少则提速数倍,多则数十倍,极端情况还有提速上千倍的。提速一个数量级基本上是常态。下面是一些案例材料:开源SPL提速保险公司团保明细单查询2000+倍开源SP... 查看详情

把numpy加速700倍?用cupy吧!

...么Numpy速度还能提升吗?本文介绍了如何利用CuPy库来加速Numpy运算速度。就其自身来说ÿ 查看详情

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

...组时,我们的3种方法分别耗时为:53.855秒0.36秒0.12秒最快加速高达448倍!那么,如何实现如此悬殊的速度优化呢?别急,实现超乎寻常的简单,Followme,你值得拥有!功能分析1. 查看详情

pythongpu加速数据科学|计算距离矩阵在用cupy时快了约100倍

文章目录一、环境配置二、计算距离矩阵CSDN叶庭云:https://yetingyun.blog.csdn.net/一、环境配置利用cuPy的cupyx.scipy.spatial.distance方法计算距离矩阵时,这个distance模块使用pylibraft作为后端,因此还需要安装好pylibraftpackage才行,可以直... 查看详情

c++-采样函数gridsampling(采样提速必备)(代码片段)

...xff0c;运用采样法提取100*100的图像数据进行分析,最少提速100倍,大大减少了不必要的计算,从而提高程序运行速度。该方法我通常用于面形拟合,来快速计算拟合系数& 查看详情

python代码加速100倍,针对excel自动化处理的加速实战!(代码片段)

并发vs并行说到并发编程,我们先来澄清一下并发(Concurrency)和并行(Parallelism)这两个概念,因为这个两个概念的含义是不同的。并行(Parallelism)指的就是在同一时刻,有两个或两个以上的任务的代码在处理器... 查看详情

pythongpu加速数据科学|计算距离矩阵在用cupy时快了约100倍

文章目录一、环境配置二、计算距离矩阵CSDN叶庭云:https://yetingyun.blog.csdn.net/一、环境配置利用cuPy的cupyx.scipy.spatial.distance方法计算距离矩阵时,这个distance模块使用pylibraft作为后端,因此还需要安装好pylibraftpackage才行,可以直... 查看详情

[codevs1243]网络提速

题目大意:有n台电脑,m个加速器,每台电脑之间传输文件有一个时间,每个加速器可以使传输时间减半(两台电脑之间可以有多个加速器),求电脑1传输文件到电脑n的最短时间。解题思路:有些人先求出最短路径,再每次找... 查看详情

用python做策略回测,耗时很长,有啥加速办法

参考技术A用Python做策略回测,耗时很长,有什么加速办法少用for,尽量用numpy/pandas的向量化方法。少用自己写的python方法,先看看numpy/pandas是不是已有现成的功能。有几个numpy的加速包,比如numexpr.安装IntelMKL.最后,可以讲关键... 查看详情

gpu加速pinterest推荐模型,参数量增加100倍,用户活跃度提高16%

作为图片届的“Twitter”,Pinterest首页展示给用户的图片也离不开背后的推荐模型。近期,其工程团队通过将机器学习服务从CPU转移到GPU上,使得Pinterest可以上线比之前大100倍的推荐模型。上线大模型给模型质量带来... 查看详情

如何让xcode在读写上提速100倍?

上个月参加了一场西雅图当地的线下iOS开发者聚会。JeffSzuhay作为一个有20+年开发经验的资深程序员,跟我讲了一套提高iOS开发效率的方法。相比于其他程序员在App启动时间、架构优化方面的经验,老爷子Jeff的优化基... 查看详情