用 Taichi 加速 Python:提速 100+ 倍!

閱讀本文大概需要 13 分鐘。

Python 已經(jīng)成為世界上最流行的編程語言，尤其在深度學(xué)習(xí)、數(shù)據(jù)科學(xué)等領(lǐng)域占據(jù)主導(dǎo)地位。但是由于其解釋執(zhí)行的屬性，Python 較低的性能很影響它在計算密集（比如多重 for 循環(huán)）的場景下發(fā)揮作用，實在讓人又愛又恨。如果你是一名經(jīng)常需要使用 Python 進(jìn)行密集計算的開發(fā)者，我相信你肯定會有下面的類似經(jīng)歷：

• 我的 Python 程序里面有個很大的 for 循環(huán)，循環(huán)體里面全是密集的計算，跑起來好慢啊...
• 我的程序里面只有一小部分計算是性能瓶頸，雖然可以用 C++ 改寫然后用 ctypes 綁定一下，但是那樣會很麻煩，還會有在別的機(jī)器上編譯不了的風(fēng)險。我希望所有的工作都能在一個 Python 腳本中完成！
• 我之前是忠實的 C++/Fortran 用戶，但是最近周圍的同學(xué)用 Python 的越來越多，我也想試試 Python，但是無奈很多祖?zhèn)鞔a用 Python 改寫以后就會慢 100 多倍，我接受不了...
• 我的工作中需要處理大量圖片數(shù)據(jù)，而需要的圖像處理功能 OpenCV 又不提供，只能自己手寫兩重 for 循環(huán)，在 Python 里面這么搞真是太痛苦了 ...

如果你有類似的煩惱，那真的值得了解一下 Taichi。我來簡單介紹一下：Taichi 是一個嵌入在 Python 中的領(lǐng)域特定語言，其一大功能就是加速 Python，讓 Python 代碼跑得和 C++ 甚至 CUDA 一樣快。Taichi 通過自己的編譯器將被 @ti.kernel 修飾的函數(shù)編譯到各種硬件上，包括 CPU 和 GPU，然后高性能執(zhí)行。

由于 Taichi 開發(fā)者社區(qū)花了大量的精力優(yōu)化 Taichi 在 Python 中的使用體驗，所有的 Taichi 功能都可以在 import taichi as ti 以后使用，Taichi 本身也可以使用 pip 進(jìn)行安裝。當(dāng)然，Taichi 也可以與常用的 Python 包（numpy、matplotlib、PyTorch 等）進(jìn)行交互。

在這篇文章中，我們將通過三個計算例子來演示如何使用 Taichi 讓你的 Python 輕松加速 > 50 倍。這三個例子是：1. 計算質(zhì)數(shù)數(shù)目；2. 動態(tài)規(guī)劃求解最長公共子序列；3. 求解反應(yīng)-擴(kuò)散方程。

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

計算素數(shù)個數(shù)

作為開胃小菜，我們先做一個小實驗：計算小于給定正整數(shù) 的素數(shù)的個數(shù)。相信任何對 Python 有基礎(chǔ)了解的人都不難寫出類似下面這樣的解法：

"""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))

這個方法的思路簡單且粗暴：我們用一個函數(shù) is_prime 來判斷某個正整數(shù) 是不是素數(shù)，是素數(shù)則返回 1，不是則返回 0。這只要遍歷檢查從 2 到 之間是否有整數(shù)能夠整除 即可。然后將小于 的全部整數(shù)依次代入此函數(shù)并統(tǒng)計結(jié)果。將上面的代碼保存為 count_primes.py，在命令行運行：

time?python?count_primes.py

在我的電腦上輸出的運行結(jié)果是：

78498

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

耗時 2.235 秒。也許代碼中 設(shè)置成一百萬對你的電腦來說太輕松了，要不要把 改成一千萬試試？我打賭不管你的電腦多么高端，你起碼都要等個半分鐘才能看到結(jié)果。

好了下面是魔法時刻：我們不修改上面的函數(shù)體，只 import 一個“庫”，然后給兩個函數(shù)分別加一個裝飾器：

"""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 命令，輸出的結(jié)果是：

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) 參數(shù)，指定 Taichi 使用 GPU 來進(jìn)行計算。在 GPU 上同樣的計算 Taichi 只花了不到 0.45 秒，比 Python 足足快了 120 倍！你可以運行這里的代碼親身體會一下。

上面這個計算素數(shù)的例子使用的方法有點土，作為習(xí)題還可以，但在實際生產(chǎn)中就顯得不那么實用了。我們接下來看一個實際中普遍使用的算法。

動態(tài)規(guī)劃

動態(tài)規(guī)劃（Dynamic Programming）是一類特別實用的算法，這類算法的哲學(xué)是以空間換時間，通過存儲中間計算結(jié)果來減少重復(fù)計算量。我們這里選擇一個求解最長公共子序列（Longest common subsequence, LCS）的例子 (算法導(dǎo)論的讀者有木有)。

插播兩個來自淵鳴的《算法導(dǎo)論》小故事：

1. 筆者小時候買過一本《算法導(dǎo)論》，書中提到四位作者都來自“麻雀理工學(xué)院”。當(dāng)時還很好奇：怎么會有學(xué)校叫這么奇怪的名字... 過了一陣才意識到自己可能成了盜版書籍的受害者。
2. 10 年后，我還真的來到了麻省理工學(xué)院（MIT）讀博士，一年后進(jìn)行碩士論文答辯（MIT 叫做 Research Qualification Exam），我自然就帶著 Taichi 的論文去了。答辯委員會里面有一位慈祥的教授，Charles E. Leiserson，嗯，就是《算法導(dǎo)論》的作者之一，“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]

最長公共子序列有很多應(yīng)用。比如大家日常使用的 Linux diff 命令和 git 工具（比較兩個文件之間的相似度），還有生物信息學(xué)中判斷兩段基因的相似度（把數(shù)字換成 ACGT 就行），其中的實現(xiàn)都用到了 LCS。

動態(tài)規(guī)劃計算 LCS 的想法是我們依次求解序列 a 的前 i 個元素和序列 b 的前 j 個元素的最長公共子序列的長度，通過讓 i 和 j 逐漸增加我們就逐步得出了最終的結(jié)果。我們用 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 的加速實現(xiàn)：

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

得到的結(jié)果為（具體結(jié)果每次未必一致）：

2721

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

我們在代碼中同時提供了分別使用 Taichi 和 Numpy 計算的版本，在我的電腦上對兩個長度是 N=15000 的隨機(jī)序列進(jìn)行計算 Taichi 版本大約需要 0.9 秒，而 Python 則需要 476s，足足差了 500 多倍！大家可以運行一下體會 Taichi 相對 Numpy 那種飛一樣的感覺。

當(dāng)然，Numpy 主要針對的場景是以數(shù)組為基本單位的運算，遇到這種需要在數(shù)組內(nèi)更細(xì)粒度進(jìn)行計算的情況就比較無力了。而這正是 Taichi 能夠發(fā)揮作用的地方。

反應(yīng) - 擴(kuò)散方程

在大自然中我們常常會在動植物的表面見到一些有趣的圖案，比如斑馬身上的條紋，獵豹身上的斑點，河豚表面的花紋等等。

這些圖案看起來是不規(guī)則的，但是又有一定的規(guī)律，并不完全隨機(jī)。從進(jìn)化的觀點，這些圖案是生物在長期演進(jìn)和自然選擇中逐漸形成的，但到底是什么規(guī)則決定了它們的形狀一直是個有趣的問題。阿蘭 . 圖靈 (正是圖靈機(jī)的發(fā)明人) 是最早注意到這一現(xiàn)象并嘗試給出模型描述的人。他在論文 "The Chemical Basis of Morphogenesis" 中提出可以用兩種化學(xué)物質(zhì) U, V 之間的相互作用來模擬圖案的形成過程，其中物質(zhì) U 的角色類似被捕食者 (prey)，物質(zhì) V 的角色類似捕食者 (predator)。它們之間的作用服從如下規(guī)則：

1. 初始時空間中隨機(jī)地分布了一些 U, V。
2. 在每個時刻 1, 2, 3, ..., U, V 兩種物質(zhì)都向其鄰域擴(kuò)散。
3. 當(dāng) U, V 相遇時，一定比例的 U, V 會合并轉(zhuǎn)化為更多的 V (捕食者在捕食后數(shù)量會增加)
4. 為了避免捕食者 V 的數(shù)量過多導(dǎo)致 U 的數(shù)量被消耗光，我們在每個時刻按照一定的比例 f 添加 U，同時按照一定的比例 k 移走 V。

于是整個過程可以用下面的反應(yīng) - 擴(kuò)散方程描述：

這里關(guān)鍵的控制參數(shù)有四個，分別是 ， 分別控制 U, V 的擴(kuò)散速度， 代表 feed，控制 U 的添加量，而 代表 kill，控制移走 V 的比例。

為了在 Taichi 中模擬這一過程，我們將空間劃分為網(wǎng)格，每個網(wǎng)格中 U, V 的濃度值用一個 vec2 來表示。注意拉普拉斯算子 的數(shù)值計算是需要訪問當(dāng)前網(wǎng)格周圍的網(wǎng)格的，為了避免一邊修改一邊讀取這種操作的發(fā)生，我們需要開辟兩個形狀為 的網(wǎng)格，每次用其中一個網(wǎng)格的值作為舊值，將更新后的濃度值寫入另一個網(wǎng)格中，然后交換兩個網(wǎng)格的角色。所以我們需要的數(shù)據(jù)結(jié)構(gòu)應(yīng)該是：

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

初始時，我們假定網(wǎng)格中 U 的濃度處處是 1，然后隨機(jī)選擇 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 的整數(shù) phase 來控制使用 uv 的哪一層來作為舊的網(wǎng)格，并將更新的值寫入 1-phase 對應(yīng)的層中。

根據(jù) V 的濃度進(jìn)行染色，我們得到了如下的動畫效果:

非常有趣的是，雖然 V 的初始濃度是隨機(jī)設(shè)置的，但是最終得到的圖案卻具有相似性。

我們在代碼中提供了基于 Taichi 和 Numba 的兩份不同的實現(xiàn)，Taichi 的版本由于使用了 GPU 進(jìn)行計算，計算的部分可以輕松達(dá)到 300+ fps，而 Numba 的版本計算部分雖然也是編譯執(zhí)行的，但由于是在 CPU 上計算的，只有大約 30fps 左右。大家可以親自運行代碼體會一下 Taichi 使用 GPU 加速的巨大優(yōu)勢。

總結(jié)

在這三個例子上 Taichi 都讓程序有了大幅加速。主要的性能來自三點：

1. Taichi 是編譯性的，而 Python 是解釋性的
2. Taichi 能自動并行，而 Python 通常是單線程的
3. Taichi 能在 GPU 上運行，而 Python 本身是在 CPU 上運行的

當(dāng)然，加速 Python 還有很多其他工具，這里我們分析一下他們和 Taichi 的優(yōu)劣。

與 Numpy/JAX/PyTorch/TensorFlow 比較：這幾類工具都高度基于數(shù)組運算。計算的最小單位是數(shù)組，在 Data Science、Deep Learning 等領(lǐng)域是有明顯的優(yōu)勢的。但是在科學(xué)計算領(lǐng)域，這樣做導(dǎo)致靈活性缺失：比如說前面那個計算質(zhì)數(shù)的程序，就比較難使用數(shù)組運算表示出來。Taichi 的優(yōu)勢就在于其靈活性，能夠直接操縱循環(huán)的每一次迭代，以一種更細(xì)顆粒度進(jìn)行對于計算的描述，類似 C++ 和 CUDA。

與 Cython 比較：使用 Cython 編寫程序?qū)崿F(xiàn)加速也是一種常見的選擇。在 Numpy 和 Scipy 的官方代碼中有不少模塊都是使用 Cython 編寫然后編譯的。但按照 Cython 的要求書寫代碼會比較麻煩，會犧牲一些可讀性。Cython 支持一定程度的并行計算，但不支持直接調(diào)用 GPU 進(jìn)行計算。

與 Numba 比較：Numba 顧名思義，是非常適合針對 Numpy 進(jìn)行加速的方案。當(dāng)你的函數(shù)是針對 Numpy 的數(shù)組向量化的操作時，使用 Numba 將其編譯以后執(zhí)行可以大大加速。Taichi 相比 Numba 的優(yōu)勢還有：1. Taichi 支持各種靈活的數(shù)據(jù)類型，比如 struct, dataclass, quant, sparse 等等，你可以任意指定它們的內(nèi)存排布，當(dāng)數(shù)據(jù)量龐大時這個優(yōu)勢會非常明顯。而 Numba 只有在針對 Numpy 的稠密數(shù)組時效果最佳。2. Taichi 可以調(diào)用不同的 GPU 后端進(jìn)行計算，所以寫大規(guī)模并行程序（如粒子仿真、渲染器等）這種操作對 Taichi 來說是小菜一碟。但你很難想象可以用 Numba 寫一個還過得去的 (哪怕離線) 渲染器。

與 Pypy 比較：Pypy 是一個 Python 的 JIT 編譯器，這個工具 2007 年就有了，和 Taichi 的解決方案有些類似，都是通過編譯的方式加速 Python。Pypy 最大優(yōu)勢在于 Python 代碼完全不用改變，就能通過 Pypy 加速。但是這也是 Pypy 加速比率比 Taichi 低的原因：因為 Pypy 需要在編譯的同時保持 Python 所有的語言特性，所以能夠進(jìn)行的優(yōu)化比較有限。而 Taichi 有一套自己的語法，雖然和 Python 很像但是也有自己的一些假設(shè)，這使得 Taichi 能夠?qū)崿F(xiàn)更大的加速。

與 ctypes 比較：ctypes 可以讓用戶在 Python 中調(diào)用 C 函數(shù)。C++、CUDA 編寫的程序也可以用過 C 接口暴露給 Python 使用。但是，ctypes 會讓工具鏈復(fù)雜化：為了寫一段比較快的程序，用戶需要同時掌握 C、Python、CMake、CUDA 等等語言，和本文描述的完全在 Python 中解決問題的方案比起來還是麻煩了一些。

總而言之，在科學(xué)計算任務(wù)上，Taichi 還是有自己獨特的優(yōu)勢的，大家可以根據(jù)自己的需求選擇最合適的工具。如果你需要在 Python 中實現(xiàn)類似 C/C++ 語言的性能，Taichi 不失為一個理想的選擇！

最后，我們希望 Taichi 能夠為你帶來價值，也希望能夠聽到你對 Taichi 的反饋，歡迎給我們提交 issues，加入 Taichi 開發(fā)者社區(qū)。如果想一鍵體驗 Taichi，只需要執(zhí)行????

pip?install?-U?taichi

并執(zhí)行????

ti?gallery

就可以體驗各種基于 Taichi 的高性能可視化 Demos，期待與大家相遇！

End

崔慶才的新書《Python3網(wǎng)絡(luò)爬蟲開發(fā)實戰(zhàn)（第二版）》已經(jīng)正式上市了！書中詳細(xì)介紹了零基礎(chǔ)用 Python 開發(fā)爬蟲的各方面知識，同時相比第一版新增了 JavaScript 逆向、Android 逆向、異步爬蟲、深度學(xué)習(xí)、Kubernetes 相關(guān)內(nèi)容，?同時本書已經(jīng)獲得 Python 之父 Guido 的推薦，目前本書正在七折促銷中！

內(nèi)容介紹：《Python3網(wǎng)絡(luò)爬蟲開發(fā)實戰(zhàn)（第二版）》內(nèi)容介紹

掃碼購買