眾所周知，Python的簡(jiǎn)單和易讀性是靠犧牲性能為代價(jià)的——

尤其是在計(jì)算密集的情況下，比如多重for循環(huán)。

不過(guò)現(xiàn)在，大佬胡淵鳴說(shuō)了：

只需import 一個(gè)叫做“Taichi”的庫(kù)，就可以把代碼速度提升100倍！

不信？

來(lái)看三個(gè)例子。

計(jì)算素?cái)?shù)的個(gè)數(shù)，速度x120

第一個(gè)例子非常非常簡(jiǎn)單，求所有小于給定正整數(shù)N的素?cái)?shù)。

標(biāo)準(zhǔn)答案如下：

我們將上面的代碼保存，運(yùn)行。

當(dāng)N為100萬(wàn)時(shí)，需要2.235s得到結(jié)果：

現(xiàn)在，我們開(kāi)始施魔法。

不用更改任何函數(shù)體，import“taichi”庫(kù)，然后再加兩個(gè)裝飾器：

Bingo！同樣的結(jié)果只要0.363s，快了將近6倍。

如果N=1000萬(wàn)，則只要0.8s；要知道，不加它可是55s，一下子又快了70倍！

不止如此，我們還可以在ti.init()中加個(gè)參數(shù)變?yōu)閠i.init(arch=ti.gpu) ，讓taich在GPU上進(jìn)行計(jì)算。

那么此時(shí)，計(jì)算所有小于1000萬(wàn)的素?cái)?shù)就只耗時(shí)0.45s了，與原來(lái)的Python代碼相比速度就提高了120倍！

厲不厲害？

什么？你覺(jué)得這個(gè)例子太簡(jiǎn)單了，說(shuō)服力不夠？我們?cè)賮?lái)看一個(gè)稍微復(fù)雜一點(diǎn)的。

動(dòng)態(tài)規(guī)劃，速度x500

動(dòng)態(tài)規(guī)劃不用多說(shuō)，作為一種優(yōu)化算法，通過(guò)動(dòng)態(tài)存儲(chǔ)中間計(jì)算結(jié)果來(lái)減少計(jì)算時(shí)間。

我們以經(jīng)典教材《算法導(dǎo)論》中的經(jīng)典動(dòng)態(tài)規(guī)劃案例“最長(zhǎng)公共子序列問(wèn)題（LCS）”為例。

比如對(duì)于序列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]。

用動(dòng)態(tài)規(guī)劃的思路計(jì)算LCS，就是先求解序列a的前i個(gè)元素和序列b的前j個(gè)元素的最長(zhǎng)公共子序列的長(zhǎng)度，然后逐步增加i或j的值，重復(fù)過(guò)程，得到結(jié)果。

我們用f[i, j]來(lái)指代這個(gè)子序列的長(zhǎng)度，即LCS((prefix(a, i), prefix(b, j)。其中prefix(a, i) 表示序列a的前i個(gè)元素，即a[0], a[1], …, a[i - 1]，得到如下遞歸關(guān)系：

完整代碼如下：

現(xiàn)在，我們用Taichi來(lái)加速：

結(jié)果如下：

胡淵鳴電腦上的程序最快做到了0.9秒內(nèi)完成，而換成用NumPy來(lái)實(shí)現(xiàn)，則需要476秒，差異達(dá)到了超500倍！

最后，我們?cè)賮?lái)一個(gè)不一樣的例子。

反應(yīng) - 擴(kuò)散方程，效果驚人

自然界中，總有一些動(dòng)物身上長(zhǎng)著一些看起來(lái)無(wú)序但實(shí)則并非完全隨機(jī)的花紋。

圖靈機(jī)的發(fā)明者艾倫·圖靈是第一個(gè)提出模型來(lái)描述這種現(xiàn)象的人。

在該模型中，兩種化學(xué)物質(zhì)（U和V）來(lái)模擬圖案的生成。這兩者之間的關(guān)系類(lèi)似于獵物和捕食者，它們自行移動(dòng)并有交互：

1. 最初，U和V隨機(jī)分布在一個(gè)域上；

2. 在每個(gè)時(shí)間步，它們逐漸擴(kuò)散到鄰近空間；

3. 當(dāng)U和V相遇時(shí)，一部分U被V吞噬。因此，V的濃度增加；

4. 為了避免U被V根除，我們?cè)诿總€(gè)時(shí)間步添加一定百分比 (f) 的U并刪除一定百分比 (k) 的V。

上面這個(gè)過(guò)程被概述為“反應(yīng)-擴(kuò)散方程”：

其中有四個(gè)關(guān)鍵參數(shù)：D_u（U的擴(kuò)散速度），D_v（V的擴(kuò)散速度），f（feed的縮寫(xiě)，控制U的加入）和k（kill的縮寫(xiě)，控制V的去除）。

如果Taichi中實(shí)現(xiàn)這個(gè)方程，首先創(chuàng)建網(wǎng)格來(lái)表示域，用vec2表示每個(gè)網(wǎng)格中U, V的濃度值。

拉普拉斯算子數(shù)值的計(jì)算需要訪問(wèn)相鄰網(wǎng)格。為了避免在同一循環(huán)中更新和讀取數(shù)據(jù)，我們應(yīng)該創(chuàng)建兩個(gè)形狀相同的網(wǎng)格W×H×2。

每次從一個(gè)網(wǎng)格訪問(wèn)數(shù)據(jù)時(shí)，我們將更新的數(shù)據(jù)寫(xiě)入另一個(gè)網(wǎng)格，然后切換下一個(gè)網(wǎng)格。那么數(shù)據(jù)結(jié)構(gòu)設(shè)計(jì)就是這樣：

一開(kāi)始，我們將U在網(wǎng)格中的濃度設(shè)置為 1，并將V放置在50個(gè)隨機(jī)選擇的位置：

那么實(shí)際計(jì)算就可以用不到10行代碼完成：

@ti.kernel
defcompute(phase:int):
fori,jinti.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

在這里，我們使用整數(shù)相位（0或1）來(lái)控制我們從哪個(gè)網(wǎng)格讀取數(shù)據(jù)。

最后一步就是根據(jù)V的濃度對(duì)結(jié)果進(jìn)行染色，就可以得到這樣一個(gè)效果驚人的圖案：?

?有趣的是，胡淵鳴介紹，即使V的初始濃度是隨機(jī)設(shè)置的，但每次都可以得到相似的結(jié)果。

而且和只能達(dá)到30fps左右的Numba實(shí)現(xiàn)比起來(lái)，Taichi實(shí)現(xiàn)由于可以選擇GPU作為后端，輕松超過(guò)了 300fps。

pip install即可安裝

看完上面三個(gè)例子，你這下相信了吧？

其實(shí)，Taichi就是一個(gè)嵌入在Python中的DSL（動(dòng)態(tài)腳本語(yǔ)言），它通過(guò)自己的編譯器將被 @ti.kernel 裝飾的函數(shù)編譯到各種硬件上，包括CPU和GPU，然后進(jìn)行高性能計(jì)算。

有了它，你無(wú)需再羨慕C++/CUDA的性能。

正如其名，Taichi就出自太極圖形胡淵鳴的團(tuán)隊(duì)，現(xiàn)在你只需要用pip install就能安裝這個(gè)庫(kù)，并與其他Python庫(kù)進(jìn)行交互，包括NumPy、Matplotlib和PyTorch等等。

當(dāng)然，Taichi用起來(lái)和這些庫(kù)以及其他加速方法有什么差別，胡淵鳴也給出了詳細(xì)的優(yōu)缺點(diǎn)對(duì)比。

審核編輯：湯梓紅