【CSDN 編者按】在本文中，我們來嘗試將 micrograd 神經(jīng)網(wǎng)絡(luò)編譯成 C。具體內(nèi)容如下：簡單了解一下神經(jīng)網(wǎng)絡(luò)；看看 micrograd 如何前向傳播和反向傳播；復(fù)習(xí)鏈?zhǔn)椒▌t；分析為什么 micrograd 的速度很慢；編寫一個(gè)小型編譯器；看看如何提高 micrograd 的速度。

原文鏈接：https://bernsteinbear.com/blog/compiling-ml-models/

未經(jīng)允許，禁止轉(zhuǎn)載！

作者 | Max Bernstein 譯者 | 彎月

責(zé)編 | 夏萌

出品 | CSDN（ID：CSDNnews）

最近，在瀏覽 Andrej Karpathy 的庫 micrograd（https://GitHub.com/karpathy/micrograd/）時(shí)，一位好友向我講解了機(jī)器學(xué)習(xí)的基礎(chǔ)知識。

micrograd 是一個(gè)純 Python 編寫的標(biāo)量值神經(jīng)網(wǎng)絡(luò)（注意計(jì)算單元不是向量，也不是矩陣），沒有用到任何庫。

micrograd 包含幾個(gè)互不相同且互補(bǔ)的部分：

• 一個(gè)基于圖的表達(dá)式生成工具和計(jì)算工具；

• 在上一步生成的計(jì)算圖上進(jìn)行反向模式自動(dòng)微分；

• 多層感知器（MLP）的神經(jīng)網(wǎng)絡(luò)構(gòu)建塊。

即便你不知道 MLP 是什么，也不必?fù)?dān)心，本文會介紹一些背景知識。如果你熟悉 Python，可以先去讀一讀 micrograd 源代碼，然后再回來繼續(xù)閱讀本文。

我們可以通過這三個(gè)主要組件編寫類似于下面的代碼：

然后就能獲得一個(gè)神經(jīng)網(wǎng)絡(luò)。

在閱讀這個(gè)代碼庫時(shí)，最初我以為這些構(gòu)建塊就是神經(jīng)網(wǎng)絡(luò)。但實(shí)際上并非如此，用建筑來做類比，這些構(gòu)建塊更像是藍(lán)圖或腳手架。每次計(jì)算神經(jīng)網(wǎng)絡(luò)時(shí)，結(jié)締組織（即中間的計(jì)算圖）都會重新構(gòu)建。用編譯器的術(shù)語來說，構(gòu)建塊有點(diǎn)像前端，而表達(dá)式圖是一種中間表示。

你可能在想我為什么要談?wù)撨@些，不是說要介紹編譯器嗎？

因?yàn)樵诶斫饬?micrograd 這三個(gè)部分之后，我意識到：

• 機(jī)器學(xué)習(xí)模型是圖；

• 前向傳播和后向傳播都是圖遍歷；

• 圖結(jié)構(gòu)不會隨時(shí)間推移而發(fā)生變化；

• 性能很重要。

這意味著，我們可以在編譯器上大做文章。這就是為什么 PyTorch 和 TensorFlow 這類的項(xiàng)目都有編譯器（TorchScript/TorchDynamo/AOT Autograd/PrimTorch/TorchInductor/Glow、XLA 等）。編譯模型可以加快訓(xùn)練和推理的速度。因此，這篇文章其實(shí)是為了通過一個(gè)很小的例子，一探大型項(xiàng)目的真實(shí)面貌。

在本文中，我們來嘗試將 micrograd 神經(jīng)網(wǎng)絡(luò)編譯成 C。具體內(nèi)容如下：

• 簡單了解一下神經(jīng)網(wǎng)絡(luò)；

• 看看 micrograd 如何前向傳播和反向傳播；

• 復(fù)習(xí)鏈?zhǔn)椒▌t；

• 分析為什么 C 的速度很慢；

• 編寫一個(gè)小型編譯器；

• 看看如何提高 micrograd 的速度。

下面，我們開始！

micrograd 實(shí)現(xiàn)的神經(jīng)網(wǎng)絡(luò)

首先，我們來了解一下多層感知器（Multi-Layer Perceptron，即MLP）。MLP 是一種密集連接的神經(jīng)網(wǎng)絡(luò)，輸入在網(wǎng)絡(luò)中沿一個(gè)方向流動(dòng)。由于上游代碼庫支持MLP，所以 micrograd 僅支持 MLP。

下面是多層感知器的示意圖：

圖：多層感知器圖。很抱歉只有一層，是我用 Excalidraw 畫的。

在此圖中，圓圈表示數(shù)據(jù)（輸入或中間計(jì)算結(jié)果），箭頭表示數(shù)據(jù)的權(quán)重和操作。在此示例中，圓圈 x、y 和 z 是輸入數(shù)據(jù)。向右的箭頭表示與權(quán)重相乘。多個(gè)箭頭指向同一個(gè)圓圈表示加法（形成點(diǎn)積），然后加上偏差（可以理解為另一個(gè)權(quán)重），所有這些都是激活函數(shù)的輸入，此示例中的激活函數(shù)為 ReLU（Rectified Linear Unit，即線性整流函數(shù)）。右邊的圓圈是第一層的結(jié)果。

Karpathy的實(shí)現(xiàn)非常直接，每個(gè)神經(jīng)元都是 Neuron 類的一個(gè)實(shí)例，并擁有一個(gè)執(zhí)行點(diǎn)積的方法 __call__。每個(gè)點(diǎn)積之后是一個(gè)激活函數(shù)，在此示例中為 ReLU，相當(dāng)于 max(x, 0)。我認(rèn)為 0 是一個(gè)隨意選取的閾值，但我不確定。

下面是 micrograd 中多層感知器的藍(lán)圖代碼（稍后我們再介紹 Value 類）：

暫時(shí)無需理會 MLP.__init__ 中使用的一些編程技巧。這確保了所有層的維度都是匹配的，同時(shí)也確保了最后一層是線性的，這意味著神經(jīng)元沒有附加激活函數(shù)。

但這個(gè)神經(jīng)網(wǎng)絡(luò)不僅僅是用浮點(diǎn)數(shù)構(gòu)建的。Karpathy 使用了 Value，為什么呢？

表達(dá)式生成器

前面我曾說過表達(dá)式圖生成器是 micrograd 的三個(gè)組件之一。

表達(dá)式生成器使用起來就像是在 Python 中通過一種稍微復(fù)雜的方法進(jìn)行數(shù)學(xué)計(jì)算：

為了方便使用，Value 類甚至實(shí)現(xiàn)了 __add__ 等所有的運(yùn)算方法，看起來與普通的 Python

數(shù)學(xué)計(jì)算非常相似。

但 Value 類不同于普通的數(shù)學(xué)計(jì)算。首先它有 grad 字段（我們稍后會詳細(xì)討論），其次它在進(jìn)行數(shù)學(xué)計(jì)算時(shí)還會構(gòu)建圖（你可以將其視為抽象語法樹（Abstract Syntax Tree，簡稱AST））。

但 Value 在正常的字符串表示形式中不可見。它的實(shí)例有一個(gè)名為 _prev 的隱藏字段，存儲了表達(dá)式的各個(gè)組成部分：

此外，Value 的實(shí)例還有一個(gè)隱藏的運(yùn)算符字段：

上述示例的意思是，名為 d 的 * 節(jié)點(diǎn)有兩個(gè)操作數(shù)：c (4) 和 a b (5)。

雖然我說過你可以把它想象成一個(gè) AST，但它不完全是 AST，因?yàn)樗皇且豢脴洹Ｋ３碛蓄愃朴谟邢驘o環(huán)圖（Directed Acyclic Graph，即DAG）的結(jié)構(gòu)。

此處，x 和 y 都使用了 w，而 z 又使用了 x 和 y，最終形成了菱形圖案。

圖：依賴關(guān)系圖，菱形依賴關(guān)系，因此是有向圖而不是樹。

假設(shè)圖中沒有環(huán)。

那么創(chuàng)建圖的代碼應(yīng)該怎么寫呢？我們應(yīng)該在 x*y 運(yùn)算的左側(cè)調(diào)用的 Value.__mul__ 函數(shù)，如下所示：

元組 children (self, other) 是指向圖中其他節(jié)點(diǎn)的指針。

但為什么我們要使用這些表達(dá)式圖呢？為什么不直接使用數(shù)學(xué)計(jì)算呢？誰會在意所有的后向指針呢？

關(guān)于梯度

訓(xùn)練神經(jīng)網(wǎng)絡(luò)其實(shí)是不斷地塑造函數(shù)（神經(jīng)網(wǎng)絡(luò)），使它能夠輸出想要的結(jié)果的過程。函數(shù)內(nèi)部有一堆系數(shù)（即權(quán)重），這些系數(shù)在訓(xùn)練過程中迭代調(diào)整。

標(biāo)準(zhǔn)的訓(xùn)練過程需要用到神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)以及另一個(gè)函數(shù)，該函數(shù)會告訴你輸出與預(yù)期值的差距（稱為損失函數(shù)）。舉一個(gè)簡單的損失函數(shù)的例子：loss（實(shí)際值, 預(yù)期值）=（預(yù)期值 – 實(shí)際值）**2，此處的 ** 代表 Python 中的求冪運(yùn)算。如果使用此函數(shù)一次處理多個(gè)輸入，則稱為“均方誤差”（MSE）。

如果你想獲得預(yù)期的輸出，則需要盡可能地最小化損失函數(shù)的值。為了最小化損失，你必須更新權(quán)重。

為了搞清楚更新哪些權(quán)重以及更新多少，你需要知道每個(gè)權(quán)重對最終損失的影響。并非每個(gè)權(quán)重的影響都是同等的，有些權(quán)重的影響更大一些。

為了計(jì)算“某個(gè)權(quán)重對損失的影響”，我們需要計(jì)算權(quán)重的梯度值（一階導(dǎo)數(shù)），即某點(diǎn)處的斜率。舉個(gè)例子，方程 y = mx b 描述的是一條直線，y 對 x 的導(dǎo)數(shù)為 m，因?yàn)?y 的值與 x 成比例，比例系數(shù)為 m，而 b 是常數(shù)。

為了計(jì)算梯度，你需要從損失值出發(fā)，向后遍歷，這個(gè)過程稱之為“反向模式自動(dòng)微分”（自動(dòng)微分，Automatic Differentiation，簡稱 AD）。聽起來很復(fù)雜，網(wǎng)上的每一篇相關(guān)文章都充斥著各種符號和曲線。但其實(shí)沒有那么難，不用害怕。

對我們來說，幸運(yùn)的是，與從上至下計(jì)算 AST 一樣，反向模式自動(dòng)微分是具有某些局部狀態(tài)的圖遍歷。如果你能編寫樹遍歷解釋器，就能實(shí)現(xiàn)反向模式自動(dòng)微分。

反向模式自動(dòng)微分與反向傳播

反向模式自動(dòng)微分并不會構(gòu)建一個(gè)對應(yīng)于普通表達(dá)式圖的導(dǎo)數(shù)圖，而是在每個(gè)節(jié)點(diǎn)的grad（梯度）字段中計(jì)算局部導(dǎo)數(shù)。然后，通過圖反向傳播這些梯度，即從損失向后一直傳播到權(quán)重。

但如何組合這些局部導(dǎo)數(shù)呢？肯定沒那么簡單吧？用數(shù)學(xué)表達(dá)式求導(dǎo)確實(shí)很復(fù)雜，但我們可以通過鏈?zhǔn)椒▌t求復(fù)合函數(shù)的導(dǎo)數(shù)。

鏈?zhǔn)椒▌t

我不是數(shù)學(xué)家。除了過去幾周重新溫習(xí)的內(nèi)容之外，我只依稀記得十年前學(xué)過鏈?zhǔn)椒▌t。如果你看不懂下面的內(nèi)容，可以通過其他方式查找詳細(xì)信息。

簡要概述

鏈?zhǔn)椒▌t告訴你如何計(jì)算復(fù)合函數(shù)的導(dǎo)數(shù)。維基百科提供了一個(gè)示例：假設(shè)函數(shù) h(x) = f(g(x))，則 h'(x) = f'(g(x)) * g'(x)（此處的 f’、h’ 和 g’ 分別是 f、h 和 g 的導(dǎo)數(shù)）。有了這條規(guī)則，組合函數(shù)就不需要任何麻煩的計(jì)算，只要了解如何獲取每個(gè)組成部分的導(dǎo)數(shù)即可。

舉個(gè)例子，假設(shè)你有 sin(x**2)，則只需知道組成函數(shù) x**2（即 2*x）和 sin(x)（即cos(x)）的導(dǎo)數(shù)，即可求出結(jié)果為：cos(x**2) * 2x。

事實(shí)證明，鏈?zhǔn)椒▌t對于大型表達(dá)式圖求導(dǎo)非常有用。你不需要仔細(xì)研究如何對龐大和過于復(fù)雜的函數(shù)求導(dǎo)。你只需要已理解的構(gòu)建塊，并且它們是組合而成的。

下面，我們將鏈?zhǔn)椒▌t應(yīng)用于表達(dá)式圖。

將鏈?zhǔn)椒▌t應(yīng)用于圖

首先，我們從一個(gè) Value 節(jié)點(diǎn)開始。對于給定的節(jié)點(diǎn)，我們可以執(zhí)行鏈?zhǔn)椒▌t的一步（偽代碼）：

此處 wrt 的意思是“對于”。求每個(gè)子節(jié)點(diǎn)對于某個(gè)子節(jié)點(diǎn)的導(dǎo)數(shù)，這一點(diǎn)非常重要。

我們不只是設(shè)置 child.grad，而是使用了 =，原因有兩個(gè)：

• 一個(gè)子節(jié)點(diǎn)可能被多個(gè)父節(jié)點(diǎn)使用，在這種情況下，子節(jié)點(diǎn)會影響到所有父節(jié)點(diǎn)。

• 批處理，暫時(shí)無需在意。

為了更具體地說明，下面我們來看看 Karpathy 實(shí)現(xiàn) * 的導(dǎo)數(shù)的方法。在數(shù)學(xué)計(jì)算中，如果 f(x,y) = x*y，則 f'(x, y) = 1*y（對于 x）且 f'(x, y) = x*1 （對于 y）。代碼如下：

這意味著，對于每個(gè)子節(jié)點(diǎn)而言，我們將使用另一個(gè)子節(jié)點(diǎn)的數(shù)據(jù)并（由于鏈?zhǔn)椒▌t）將其與父表達(dá)式的梯度相乘。也就是說，self.grad（左側(cè)）是使用 other.data（右側(cè)）調(diào)整的，反之亦然。我們通過上述步驟成功地將數(shù)學(xué)計(jì)算轉(zhuǎn)換成了代碼。求導(dǎo)數(shù)，應(yīng)用鏈?zhǔn)椒▌t，然后加到子節(jié)點(diǎn)的梯度上。

到這里，我們建立了一個(gè)函數(shù)來求操作節(jié)點(diǎn)的導(dǎo)數(shù)，但我們需要對整個(gè)圖求導(dǎo)。

遍歷圖并不像遍歷樹那么簡單。你需要避免重復(fù)訪問同一個(gè)節(jié)點(diǎn)，并保證在父節(jié)點(diǎn)之前訪問子節(jié)點(diǎn)（正向模式）或在子節(jié)點(diǎn)之前訪問父節(jié)點(diǎn)（反向模式）。難點(diǎn)在于，雖然我們不會重復(fù)訪問同一個(gè)節(jié)點(diǎn)，但訪問會更新該節(jié)點(diǎn)的子節(jié)點(diǎn)（而不是節(jié)點(diǎn)本身），并且多個(gè)節(jié)點(diǎn)可能共享子節(jié)點(diǎn)，因此子節(jié)點(diǎn)的梯度可能會被多次更新。這都是正常情況。

為此，我們必須引入拓?fù)渑判?/span>。

拓?fù)渑判蚝蛨D轉(zhuǎn)換

圖的拓?fù)渑判蚰鼙ＷC圖中的子節(jié)點(diǎn)永遠(yuǎn)先于父節(jié)點(diǎn)被訪問。一般來說，只有當(dāng)圖中沒有環(huán)路時(shí)才能使用拓?fù)渑判?，幸運(yùn)的是，我們前面已經(jīng)假設(shè)圖沒有環(huán)路。

下面是 Value 圖的拓?fù)渑判蚴纠榱撕啙嵠鹨?，我們使用了嵌套函?shù) build_topo，但這并不是絕對必要的。

為了說明上述代碼的工作原理，我們可以針對一個(gè)非常簡單的表達(dá)式圖 1*2 進(jìn)行拓?fù)渑判颉?/p>

在這段拓?fù)渑判蛑?，為了?jì)算值 3，首先我們必須計(jì)算值 1 和 2。值 1 和 2 的計(jì)算順序并不重要，但二者都必須在 3 之前計(jì)算。

以上，我們確定了圖的遍歷順序，下面可以著手反向傳播了。

將拓?fù)渑判驊?yīng)用于反向傳播

我們可以利用上述介紹的鏈?zhǔn)椒▌t和拓?fù)渑判?，在圖上進(jìn)行反向傳播。下面是 micrograd 的實(shí)現(xiàn)代碼。首先構(gòu)建一個(gè)拓?fù)渑判?，然后對其進(jìn)行反向操作，將鏈?zhǔn)椒▌t應(yīng)用于每個(gè) Value。

通常，我們會在損失函數(shù)的結(jié)果 Value 上調(diào)用 Value.backward。

你可能在想為什么我們在進(jìn)行反向傳播之前將 self.grad 設(shè)置為 1，你可以仔細(xì)想一想。

整合

此處我不打算深入細(xì)節(jié)，只是粗略地介紹一下如何通過簡單的訓(xùn)練，使用基于 MLP 的分類器解決 MNIST 數(shù)字識別問題。下面的代碼不能直接運(yùn)行，其中缺少圖像加載支持代碼和損失函數(shù)。超參數(shù)（批量大小等）是任意設(shè)定的，且未經(jīng)調(diào)整。完整的訓(xùn)練代碼和添加 exp/log/Max 的引擎改動(dòng)，請參見 GitHub 代碼庫。

• 訓(xùn)練代碼：https://github.com/tekknolagi/micrograd/blob/534ab3c884e66c8a325e0a8f3ed278656a616002/mnist.py

• 相應(yīng)的引擎改動(dòng)：https://github.com/tekknolagi/micrograd/blob/534ab3c884e66c8a325e0a8f3ed278656a616002/micrograd/engine.py

在上述代碼片段中，MLP (model = MLP(…)) 會在各層中構(gòu)建一堆神經(jīng)元，并將一些權(quán)重初始化為 Value 的實(shí)例，但它還沒有構(gòu)建圖。只有被調(diào)用時(shí)（model(image.pixels)），它才會構(gòu)建圖并執(zhí)行所有點(diǎn)積。然后，在計(jì)算損失時(shí)，我們在此基礎(chǔ)上構(gòu)建更多的圖。這就是前向傳播。

接著，我們反向傳播，利用損失調(diào)用 backward。

然后，我們通過梯度來調(diào)整所有權(quán)重。

最后，請不要忘記將梯度初始化為零！

性能問題

用 CPython 運(yùn)行這段代碼會非常慢，感覺上計(jì)算一張圖的前向傳播大約需要一秒鐘，我們還需要反向傳播，而且我們必須對 6萬 張圖進(jìn)行幾個(gè) epoch 的處理。最終花費(fèi)的時(shí)間會過長。

我們可以按照大家的建議，嘗試使用 PyPy。這樣每秒可以處理幾張圖，但仍然不夠快。

順便說一下，我們的舊項(xiàng)目 Skybison 比 CPython 和 PyPy 都快得多！經(jīng)過分析后，我們發(fā)現(xiàn)性能的主要痛點(diǎn)是函數(shù)創(chuàng)建（在 Skybison中有點(diǎn)慢），但如果將內(nèi)部函數(shù) _backward 放到頂層，問題就會消失。因此，很明顯拓?fù)渑判虻募喜檎沂桥渲梦募凶盥牟糠?。之后是所有臨時(shí) Value 對象的垃圾回收。

另外，將內(nèi)部函數(shù)放到頂層也可以極大地提高 PyPy 的速度，并且比 Skybison 更快。

我認(rèn)為所有運(yùn)行時(shí)的痛點(diǎn)是：

• 每次前向傳遞都會重新創(chuàng)建圖，因?yàn)樗械?Value及其 _backward 函數(shù)必須重新分配。

  Neuron.__call__ 中的 zip 也存在大量內(nèi)存分配和迭代開銷。

• 由于指針追逐、函數(shù)調(diào)用以及集合/列表內(nèi)存分配和操作，每次反向傳播都會進(jìn)行拓?fù)渑判颉?/p>

• 正常的 Python解釋器開銷。

但根據(jù)多年的經(jīng)驗(yàn)，我認(rèn)為首先應(yīng)該實(shí)際測量一下，而不是在黑暗中盲目優(yōu)化。

使用分析器檢查

Emery Berger 和他的團(tuán)隊(duì)發(fā)布了一款出色的 Python 分析工具，名為 Scalene。使用時(shí)，你可以直接運(yùn)行 scalene yourprogram.py（代替python3 yourprogram.py）。程序運(yùn)行完成后（或者按Control-C鍵），將彈出一個(gè)本地托管的小型網(wǎng)站，其中包含分析信息。

我在 micrograd MNIST 上運(yùn)行了 Scalene，結(jié)果如下：

圖：Scalene 分析器輸出的 micrograd 分析結(jié)果屏幕截圖。

我們可以看到許多 Value 的內(nèi)存分配，而且 self._prev 是一個(gè)集合，甚至有可能造成內(nèi)存泄漏。特別是，我們還可以看到很多 和 * 操作，因?yàn)?__add__和 __mul__ 分配了很多內(nèi)存。

看看內(nèi)存使用列，曲線呈向上向右的趨勢，這不是我們想要的。似乎為每個(gè) Value 創(chuàng)建 _prev 元素的set 的過程占據(jù)了大量時(shí)間。

如果你是守舊派，不信任新的分析工具，那么甚至可以使用 perf 確認(rèn)這些觀察結(jié)果。你可能需要根據(jù) Python 發(fā)行版安裝調(diào)試符號（我使用的是 Ubuntu 的 python3.10-dbg），然后運(yùn)行 perf record python3 yourprogram.py。我得到的結(jié)果如下（省略了 0.5% 以下的記錄）：

gc_collect_main占總體時(shí)間的 37% 是一個(gè)巨大的危險(xiǎn)信號。其次，下面的其他函數(shù)（deduce_unreachable和所有的 _traverse 函數(shù)）看起來也與垃圾回收相關(guān)，這意味著程序占用了太多內(nèi)存。所以，Scalene 和 perf 的分析結(jié)果似乎是一致的。

如果去掉 set(_children)，僅使用元組（這似乎不會影響正確性），則時(shí)間占用會相對分散一些。

還有一個(gè)簡單的方法是將 __slots__ 添加到 Value 類。屬性字典是我能想到的唯一分配字典的地方，所以也許我們可以解決這個(gè)問題。果然 添加 __slots__ 后，dict_traverse 就消失了。

最后，我們還可以嘗試刪除嵌套函數(shù)分配，這樣就可以消除 func_traverse。不過這項(xiàng)優(yōu)化要比前兩個(gè)更加繁瑣一點(diǎn)。

這些小改動(dòng)不會改變程序的整體架構(gòu)，所以也不會帶來大量的數(shù)學(xué)運(yùn)算和圖遍歷工作。

那么，我們應(yīng)該怎么辦呢？

解決方案

提高程序運(yùn)行速度的最佳方法是減少工作量。垃圾回收太多？則減少內(nèi)存分配。遞歸太多？則減少拓?fù)?/span>分類。開銷太大？則減少解釋的工作量。更詳細(xì)地說，我提出的解決方案是：

• 重用輸入之間的圖結(jié)構(gòu)。避免每次都重新構(gòu)建 Value 圖，復(fù)制新輸入并執(zhí)行前向傳播和反向傳播。

• 由于不修改圖，因此也無需重新拓?fù)渑判?。順序保持不變，有利于前向傳播和反向傳播?/p>

• 歸根結(jié)底，Value抽象并不重要。如果我們知道遍歷的順序并使用 IEEE-754 雙精度，就應(yīng)該將拓?fù)渑判蚣捌洳僮骶幾g為 C 或更簡單的東西。

這與我們了解的編譯器知識相一致：如果可以在程序允許的語義中凍結(jié)一些動(dòng)態(tài)，就可以提升性能。由于圖是靜態(tài)的，所以我們可以采用這種做法。

編寫編譯器

這個(gè)編譯器的目標(biāo)是，為 micrograd 編寫一款非常小且非常適合的編譯器，不需要重新設(shè)計(jì)架構(gòu)。

我們可以編寫一種字節(jié)碼編譯器，并通過這個(gè)編譯器去除所有函數(shù)調(diào)用以及重復(fù)的樹遍歷和指針追逐。這樣就能提升性能。但不幸的是，我們?nèi)匀挥幸粋€(gè)解釋器循環(huán)，并且該解釋器是用 Python 編寫的，這會產(chǎn)生大量開銷。

為此，我們將進(jìn)一步將這些代碼編譯為 C。最終目標(biāo)是編寫一個(gè) Python C 插件，我們可以導(dǎo)入并使用它來代替 micrograd 的解釋版本。

該項(xiàng)目的最初版本是將 MLP、Layer 和 Neuron 類直接編譯為 C，但不幸的是，它的可擴(kuò)展性不是很好：修改模型的架構(gòu)需要編寫新的編譯器。此外， 它也不支持反向傳播，只是對推理有幫助。

因此，我們需要為 Value 圖編寫編譯器。這意味著，只要機(jī)器學(xué)習(xí)架構(gòu)使用 Values，那么任何人都可以毫不費(fèi)力地使用這款編譯器。你只需要為它寫一個(gè)解釋器。

前向傳播

由于我們使用了拓?fù)渑判?，所以不妨在前向傳播和反向傳播中使用它。那么，我們只需要編寫一個(gè)一次只處理一個(gè) Value 的編譯器。寫法如下：

（假設(shè) data 是我們稍后將創(chuàng)建的一個(gè)大小適當(dāng)?shù)碾p精度數(shù)組。）

上面的代碼把圖變成了線性。這有點(diǎn)像我們之前看到的拓?fù)渑判颍鞘褂?C 代碼編寫的。這種策略之所以有效，是因?yàn)槲覀儧]有循環(huán)，也沒有重新定義 Values。每個(gè)值設(shè)置一次，而且這段代碼即使包含內(nèi)存加載和存儲，也應(yīng)該比 Python 中的指針追蹤和函數(shù)調(diào)用快得多。

我們可以編寫一個(gè)類似的解釋版本，每種操作都有自己的方法（__add__、__mul__ 等），但將編譯器全部呈現(xiàn)在一個(gè)方法中更容易。因此我添加了一個(gè)編譯函數(shù)。下面的示例實(shí)現(xiàn)了常量值（op==”）和加法（op==’ ‘）：

其他運(yùn)算符的寫法也一樣。例如，你可以試試看如何實(shí)現(xiàn) ** 或 exp。請注意，** 需要存儲額外的數(shù)據(jù)或某種特殊的處理。

你可能會注意到，這種編譯策略需要為 Values 分配 ID。為此，我在__init__ 函數(shù)中添加了一個(gè) _id 字段，它是__init__ 函數(shù)中的自動(dòng)遞增計(jì)數(shù)器。具體的實(shí)現(xiàn)并不重要，你只需要知道每個(gè) Value 對象都有一個(gè)唯一的 _id 即可。

我的編譯器實(shí)現(xiàn)所有的操作大約為 40 行代碼，甚至包括一些小的即時(shí)優(yōu)化。但這個(gè)編譯器是前向傳播。反向傳播呢？我們也需要加快訓(xùn)練速度。向后傳播一定更為復(fù)雜，是嗎？

反向傳播

實(shí)際上，反向傳播的復(fù)雜度與前向傳播差不多。我們只需要逐行修改反向傳播函數(shù)（所有的 _backward 實(shí)現(xiàn)）。

例如，我們需要修改 * 的反向傳播。我添加了一些輔助函數(shù)，這樣代碼行數(shù)更少，而且看起來更像解釋版本。與前向傳播一樣，所有運(yùn)算符都在一個(gè)方法：backward_compile。

（與前向傳播一樣，我們假設(shè) grad 是稍后即將創(chuàng)建的一個(gè)大小適中的雙精度數(shù)組。）

下面，我們來看看如何應(yīng)用：

很奇怪，為什么 x (grad[6]) 和 y (grad[7]) 沒有反向傳播代碼？因?yàn)樗鼈儧]有子節(jié)點(diǎn)，而本身由父節(jié)點(diǎn) z (grad[8]) 調(diào)整。我前面曾提到過，訪問節(jié)點(diǎn)會調(diào)整該節(jié)點(diǎn)的子節(jié)點(diǎn)。

我的編譯器實(shí)現(xiàn)反向傳播大約為 30 行代碼，甚至比前向傳播還短，非常整潔。

如此，我們就完成了編譯器的編寫。恭喜！本文最復(fù)雜的部分已經(jīng)結(jié)束了。其余都是一些小細(xì)節(jié)和 Python C-API 的具體實(shí)現(xiàn)。

更新權(quán)重

在實(shí)現(xiàn)了反向傳播后，我們需要通過它們的梯度來調(diào)整權(quán)重。此處的代碼只需將 Python 代碼機(jī)械地翻譯成 C。為了方便比較，下面是解釋版本：

它會在運(yùn)行期間遍歷并調(diào)整模型參數(shù)。相比之下，編譯版本在編譯時(shí)進(jìn)行迭代，而在運(yùn)行時(shí)僅執(zhí)行減法操作：

如果不考慮 assert，上述代碼的長度與 Python 相差無幾。

設(shè)置輸入

當(dāng)輸入不是整數(shù)和浮點(diǎn)數(shù)等簡單數(shù)據(jù)類型時(shí)，將 Python 代碼輸入到 C 會有點(diǎn)棘手。理想情況下，我們生成的機(jī)器學(xué)習(xí)代碼能夠與 Python 共享內(nèi)存，這樣就可以避免來回復(fù)制數(shù)據(jù)，但這種實(shí)現(xiàn)并不簡單，因此我們需要采用略麻煩的做法。

我們來添加一個(gè)函數(shù) set_input，將黑白像素?cái)?shù)據(jù)放入字節(jié)數(shù)組中，并將每個(gè)像素復(fù)制到 data 數(shù)組的每個(gè)槽中。雖然這種方法相對較慢，但肯定不會成為管道中的瓶頸。

在這個(gè)例子中，inp 是輸入數(shù)組。與 micrograd 的解釋版本不同，我們不會在每次迭代中創(chuàng)建新的輸入 Values。這意味著，我們必須預(yù)先分配機(jī)器學(xué)習(xí)模型輸入和輸出的 ID 范圍：

請注意，由于 inp 和 exp 是任意選擇的，所以每個(gè) Value 節(jié)點(diǎn)的 data 或grad 字段都包含垃圾數(shù)據(jù)。但是，生成的 C 代碼并不會使用這些 Python 值。我們關(guān)心的是 _op 和 _prev 字段表示的圖結(jié)構(gòu)。

為了使用 Python 中的 C 代碼，我們必須使用 C-API 來創(chuàng)建 Python C 插件。

Python C 插件

通過一堆代碼更新 data 和 grad 度數(shù)組很有趣，而且它是一個(gè)完整的編譯器，但目前還不能發(fā)揮作用。我們需要將這些代碼包裝到函數(shù)中（我為它們?nèi)∶麨?forward、backward、update和 set_input），并允許 Python 驅(qū)動(dòng)程序訪問這些函數(shù)。我們不想完全使用 C！

大部分的代碼很簡單（只需要添加 print(“void forward() {” 等代碼），但部分代碼需要用到 Python 的內(nèi)部知識。

例如，下面是包裝 forward 函數(shù)的代碼：

這是一個(gè)fastcall C-API 函數(shù)的示例，這意味著它會通過一個(gè)數(shù)組獲取參數(shù)。我們必須像下面這樣注冊這個(gè)函數(shù)：

接下來，我們創(chuàng)建一個(gè)可供 Python 導(dǎo)入的模塊描述，這樣就可以在導(dǎo)入時(shí)創(chuàng)建模塊對象：

然后，我們來創(chuàng)建 PyInit_nn 函數(shù)。如果 Python 的原生導(dǎo)入器在.so 中找到模塊，并且這個(gè)模塊擁有 PyInit_XYZ函數(shù)，則調(diào)用它來創(chuàng)建模塊對象。

到此，我們的編譯器就基本編寫完成了。接下來的主要工作是模型的訓(xùn)練和推理。

正確嗎？速度提升了嗎？

這是兩個(gè)不同的問題，如果代碼產(chǎn)生錯(cuò)誤的輸出，那么性能就沒有任何意義了。

正確性

測試編譯器有些棘手。編譯器的很多部分不僅必須能夠單獨(dú)工作，同時(shí)必須與其他部分協(xié)同工作。值得慶幸的是，在這個(gè)示例中，我們的編譯器非常小，只包含少量基本操作。因此，為生成的 C 代碼編寫單元測試并不太難。

此外，我們還可以針對同一段代碼的解釋版本和編譯版本輸出的數(shù)字進(jìn)行一些并行測試。如果二者都有一定的誤差范圍，則我們可以認(rèn)為編譯器是正確的。不過，我不建議使用 MNIST。解釋版本太慢，而單元測試應(yīng)該能夠很快地運(yùn)行?；蛟S可以試試看XOR。

值得慶幸的是，CPython 的 float 使用了宿主系統(tǒng)的浮點(diǎn)實(shí)現(xiàn)，因此我們無需額外的努力即可獲得與 C 相同的數(shù)字行為。

性能

在我的機(jī)器上，訓(xùn)練從每秒 1 個(gè)圖像（解釋版本）上升到了每秒 > 1000 個(gè)圖像（編譯版本），性能提升大約為 1 千倍！不過，編譯版本會產(chǎn)生一些前期的開銷，因?yàn)槟惚仨毦幾g C 代碼。如果使用 TCC（一種速度非?？斓?C 編譯器），那么可以獲得非常好的性能。我的編譯時(shí)間大約為半秒，每個(gè) epoch 大約需要 45 秒。如果使用 Clang（一種相對很慢的 C 編譯器），則可以獲得更好的性能。具體的數(shù)字如下表所示：

編譯時(shí)間（秒）	每個(gè) epoch 所需時(shí)間（秒）	速度提升
解釋版本	0	60,000	1倍
TCC	0.5	45	1333倍
Clang -O0	~30	30	2000倍
Clang -O1	~350	8	7500倍

不管怎樣看，這都是一個(gè)巨大的勝利。我認(rèn)為我們成功了！

完整的編譯器代碼、編譯器包裝器和訓(xùn)練代碼，請參見 GitHub：

• 編譯器代碼：https://github.com/tekknolagi/micrograd/blob/c15b6b8fd373c48014be369c4f7bd0917932a53b/micrograd/engine.py

• 編譯器包裝器和訓(xùn)練代碼：https://github.com/tekknolagi/micrograd/blob/c15b6b8fd373c48014be369c4f7bd0917932a53b/test.py

總結(jié)

神經(jīng)網(wǎng)絡(luò)由靜態(tài)數(shù)據(jù)流圖表示，可向前或向后執(zhí)行。這意味著，它們有點(diǎn)像遍歷樹的解釋器。這也意味著，將樹編譯為較低級別的表示可以加快程序的運(yùn)行速度。

歡迎參與 CSDN 重磅發(fā)起的《2023 AI 開發(fā)者生態(tài)調(diào)查問卷》，分享您真實(shí)的 AI 使用體驗(yàn)，更有精美好禮等你拿！