【CSDN 編者按】在本文中，我們來嘗試將 micrograd 神經(jīng)網(wǎng)絡(luò)編譯成 C。具體內(nèi)容如下：簡單了解一下神經(jīng)網(wǎng)絡(luò)；看看 micrograd 如何前向傳播和反向傳播；復(fù)習(xí)鏈?zhǔn)椒▌t；分析為什么 micrograd 的速度很慢；編寫一個小型編譯器；看看如何提高 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/）時，一位好友向我講解了機器學(xué)習(xí)的基礎(chǔ)知識。

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

micrograd 包含幾個互不相同且互補的部分：

• 一個基于圖的表達式生成工具和計算工具；

• 在上一步生成的計算圖上進行反向模式自動微分；

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

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

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

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

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

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

因為在理解了 micrograd 這三個部分之后，我意識到：

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

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

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

• 性能很重要。

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

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

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

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

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

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

• 編寫一個小型編譯器；

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

下面，我們開始！

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

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

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

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

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

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

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

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

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

表達式生成器

前面我曾說過表達式圖生成器是 micrograd 的三個組件之一。

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

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

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

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

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

此外，Value 的實例還有一個隱藏的運算符字段：

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

雖然我說過你可以把它想象成一個 AST，但它不完全是 AST，因為它不是一棵樹。它常常擁有類似于有向無環(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 運算的左側(cè)調(diào)用的 Value.__mul__ 函數(shù)，如下所示：

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

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

關(guān)于梯度

訓(xùn)練神經(jīng)網(wǎng)絡(luò)其實是不斷地塑造函數(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)以及另一個函數(shù)，該函數(shù)會告訴你輸出與預(yù)期值的差距（稱為損失函數(shù)）。舉一個簡單的損失函數(shù)的例子：loss（實際值, 預(yù)期值）=（預(yù)期值 – 實際值）**2，此處的 ** 代表 Python 中的求冪運算。如果使用此函數(shù)一次處理多個輸入，則稱為“均方誤差”（MSE）。

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

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

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

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

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

反向模式自動微分與反向傳播

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

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

鏈?zhǔn)椒▌t

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

簡要概述

鏈?zhǔn)椒▌t告訴你如何計算復(fù)合函數(shù)的導(dǎo)數(shù)。維基百科提供了一個示例：假設(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ù)就不需要任何麻煩的計算，只要了解如何獲取每個組成部分的導(dǎo)數(shù)即可。

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

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

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

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

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

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

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

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

• 批處理，暫時無需在意。

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

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

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

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

為此，我們必須引入拓撲排序。

拓撲排序和圖轉(zhuǎn)換

圖的拓撲排序能保證圖中的子節(jié)點永遠先于父節(jié)點被訪問。一般來說，只有當(dāng)圖中沒有環(huán)路時才能使用拓撲排序，幸運的是，我們前面已經(jīng)假設(shè)圖沒有環(huán)路。

下面是 Value 圖的拓撲排序示例。為了簡潔起見，我們使用了嵌套函數(shù) build_topo，但這并不是絕對必要的。

為了說明上述代碼的工作原理，我們可以針對一個非常簡單的表達式圖 1*2 進行拓撲排序。

在這段拓撲排序中，為了計算值 3，首先我們必須計算值 1 和 2。值 1 和 2 的計算順序并不重要，但二者都必須在 3 之前計算。

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

將拓撲排序應(yīng)用于反向傳播

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

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

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

整合

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

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

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

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

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

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

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

性能問題

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

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

順便說一下，我們的舊項目 Skybison 比 CPython 和 PyPy 都快得多！經(jīng)過分析后，我們發(fā)現(xiàn)性能的主要痛點是函數(shù)創(chuàng)建（在 Skybison中有點慢），但如果將內(nèi)部函數(shù) _backward 放到頂層，問題就會消失。因此，很明顯拓撲排序的集合查找是配置文件中最慢的部分。之后是所有臨時 Value 對象的垃圾回收。

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

我認為所有運行時的痛點是：

• 每次前向傳遞都會重新創(chuàng)建圖，因為所有的 Value及其 _backward 函數(shù)必須重新分配。

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

• 由于指針追逐、函數(shù)調(diào)用以及集合/列表內(nèi)存分配和操作，每次反向傳播都會進行拓撲排序。

• 正常的 Python解釋器開銷。

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

使用分析器檢查

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

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

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

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

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

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

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

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

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

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

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

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

解決方案

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

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

• 由于不修改圖，因此也無需重新拓撲排序。順序保持不變，有利于前向傳播和反向傳播。

• 歸根結(jié)底，Value抽象并不重要。如果我們知道遍歷的順序并使用 IEEE-754 雙精度，就應(yīng)該將拓撲排序及其操作編譯為 C 或更簡單的東西。

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

編寫編譯器

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

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

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

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

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

前向傳播

由于我們使用了拓撲排序，所以不妨在前向傳播和反向傳播中使用它。那么，我們只需要編寫一個一次只處理一個 Value 的編譯器。寫法如下：

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

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

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

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

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

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

反向傳播

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

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

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

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

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

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

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

更新權(quán)重

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

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

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

設(shè)置輸入

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

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

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

請注意，由于 inp 和 exp 是任意選擇的，所以每個 Value 節(jié)點的 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ù)組很有趣，而且它是一個完整的編譯器，但目前還不能發(fā)揮作用。我們需要將這些代碼包裝到函數(shù)中（我為它們?nèi)∶麨?forward、backward、update和 set_input），并允許 Python 驅(qū)動程序訪問這些函數(shù)。我們不想完全使用 C！

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

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

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

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

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

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

正確嗎？速度提升了嗎？

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

正確性

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

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

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

性能

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

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

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

完整的編譯器代碼、編譯器包裝器和訓(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í)行。這意味著，它們有點像遍歷樹的解釋器。這也意味著，將樹編譯為較低級別的表示可以加快程序的運行速度。

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