我把 ML 模型编译成 C 后，速度竟提升了 1000 倍！

作者 | Max Bernstein       译者 | 弯月

责编 | 夏萌

出品 | CSDN（ID：CSDNnews）

最近，在浏览 Andrej Karpathy 的库 micrograd（https://github.com/karpathy/micrograd/）时，一位好友向我讲解了机器学习的基础知识。

micrograd 是一个纯 Python 编写的标量值神经网络（注意计算单元不是向量，也不是矩阵），没有用到任何库。

micrograd 包含几个互不相同且互补的部分：

一个基于图的表达式生成工具和计算工具；

在上一步生成的计算图上进行反向模式自动微分；

多层感知器（MLP）的神经网络构建块。

即便你不知道 MLP 是什么，也不必担心，本文会介绍一些背景知识。如果你熟悉 Python，可以先去读一读 micrograd 源代码，然后再回来继续阅读本文。

我们可以通过这三个主要组件编写类似于下面的代码：

然后就能获得一个神经网络。

在阅读这个代码库时，最初我以为这些构建块就是神经网络。但实际上并非如此，用建筑来做类比，这些构建块更像是蓝图或脚手架。每次计算神经网络时，结缔组织（即中间的计算图）都会重新构建。用编译器的术语来说，构建块有点像前端，而表达式图是一种中间表示。

你可能在想我为什么要谈论这些，不是说要介绍编译器吗？

因为在理解了 micrograd 这三个部分之后，我意识到：

机器学习模型是图；

前向传播和后向传播都是图遍历；

图结构不会随时间推移而发生变化；

性能很重要。

这意味着，我们可以在编译器上大做文章。这就是为什么 PyTorch 和 TensorFlow 这类的项目都有编译器（TorchScript/TorchDynamo/AOT Autograd/PrimTorch/TorchInductor/Glow、XLA 等）。编译模型可以加快训练和推理的速度。因此，这篇文章其实是为了通过一个很小的例子，一探大型项目的真实面貌。

在本文中，我们来尝试将 micrograd 神经网络编译成 C。具体内容如下：

简单了解一下神经网络；

看看 micrograd 如何前向传播和反向传播；

复习链式法则；

分析为什么 C 的速度很慢；

编写一个小型编译器；

看看如何提高 micrograd 的速度。

下面，我们开始！

micrograd 实现的神经网络

首先，我们来了解一下多层感知器（Multi-Layer Perceptron，即MLP）。MLP 是一种密集连接的神经网络，输入在网络中沿一个方向流动。由于上游代码库支持MLP，所以 micrograd 仅支持 MLP。

下面是多层感知器的示意图：

图：多层感知器图。很抱歉只有一层，是我用 Excalidraw 画的。

在此图中，圆圈表示数据（输入或中间计算结果），箭头表示数据的权重和操作。在此示例中，圆圈 x、y 和 z 是输入数据。向右的箭头表示与权重相乘。多个箭头指向同一个圆圈表示加法（形成点积），然后加上偏差（可以理解为另一个权重），所有这些都是激活函数的输入，此示例中的激活函数为 ReLU（Rectified Linear Unit，即线性整流函数）。右边的圆圈是第一层的结果。

Karpathy的实现非常直接，每个神经元都是 Neuron 类的一个实例，并拥有一个执行点积的方法 __call__。每个点积之后是一个激活函数，在此示例中为 ReLU，相当于 max(x, 0)。我认为 0 是一个随意选取的阈值，但我不确定。

下面是 micrograd 中多层感知器的蓝图代码（稍后我们再介绍 Value 类）：

暂时无需理会 MLP.__init__ 中使用的一些编程技巧。这确保了所有层的维度都是匹配的，同时也确保了最后一层是线性的，这意味着神经元没有附加激活函数。

但这个神经网络不仅仅是用浮点数构建的。Karpathy 使用了 Value，为什么呢？

表达式生成器

前面我曾说过表达式图生成器是 micrograd 的三个组件之一。

表达式生成器使用起来就像是在 Python 中通过一种稍微复杂的方法进行数学计算：

为了方便使用，Value 类甚至实现了 __add__ 等所有的运算方法，看起来与普通的 Python

数学计算非常相似。

但 Value 类不同于普通的数学计算。首先它有 grad 字段（我们稍后会详细讨论），其次它在进行数学计算时还会构建图（你可以将其视为抽象语法树（Abstract Syntax Tree，简称AST））。

但 Value 在正常的字符串表示形式中不可见。它的实例有一个名为 _prev 的隐藏字段，存储了表达式的各个组成部分：

此外，Value 的实例还有一个隐藏的运算符字段：

上述示例的意思是，名为 d 的 * 节点有两个操作数：c(4) 和 a + b (5)。

虽然我说过你可以把它想象成一个 AST，但它不完全是 AST，因为它不是一棵树。它常常拥有类似于有向无环图（Directed Acyclic Graph，即DAG）的结构。

此处，x 和 y 都使用了 w，而 z 又使用了 x 和 y，最终形成了菱形图案。

图：依赖关系图，菱形依赖关系，因此是有向图而不是树。

假设图中没有环。

那么创建图的代码应该怎么写呢？我们应该在 x*y 运算的左侧调用的 Value.__mul__ 函数，如下所示：

元组 children (self, other) 是指向图中其他节点的指针。

但为什么我们要使用这些表达式图呢？为什么不直接使用数学计算呢？谁会在意所有的后向指针呢？

关于梯度

训练神经网络其实是不断地塑造函数（神经网络），使它能够输出想要的结果的过程。函数内部有一堆系数（即权重），这些系数在训练过程中迭代调整。

标准的训练过程需要用到神经网络结构以及另一个函数，该函数会告诉你输出与预期值的差距（称为损失函数）。举一个简单的损失函数的例子：loss（实际值, 预期值）=（预期值 - 实际值）**2，此处的 ** 代表 Python 中的求幂运算。如果使用此函数一次处理多个输入，则称为“均方误差”（MSE）。

如果你想获得预期的输出，则需要尽可能地最小化损失函数的值。为了最小化损失，你必须更新权重。

为了搞清楚更新哪些权重以及更新多少，你需要知道每个权重对最终损失的影响。并非每个权重的影响都是同等的，有些权重的影响更大一些。

为了计算“某个权重对损失的影响”，我们需要计算权重的梯度值（一阶导数），即某点处的斜率。举个例子，方程 y = mx + b 描述的是一条直线，y 对 x 的导数为 m，因为 y 的值与 x 成比例，比例系数为 m，而 b 是常数。

为了计算梯度，你需要从损失值出发，向后遍历，这个过程称之为“反向模式自动微分”（自动微分，AutomaticDifferentiation，简称 AD）。听起来很复杂，网上的每一篇相关文章都充斥着各种符号和曲线。但其实没有那么难，不用害怕。

对我们来说，幸运的是，与从上至下计算 AST 一样，反向模式自动微分是具有某些局部状态的图遍历。如果你能编写树遍历解释器，就能实现反向模式自动微分。

反向模式自动微分与反向传播

反向模式自动微分并不会构建一个对应于普通表达式图的导数图，而是在每个节点的grad（梯度）字段中计算局部导数。然后，通过图反向传播这些梯度，即从损失向后一直传播到权重。

但如何组合这些局部导数呢？肯定没那么简单吧？用数学表达式求导确实很复杂，但我们可以通过链式法则求复合函数的导数。

链式法则

我不是数学家。除了过去几周重新温习的内容之外，我只依稀记得十年前学过链式法则。如果你看不懂下面的内容，可以通过其他方式查找详细信息。

简要概述

链式法则告诉你如何计算复合函数的导数。维基百科提供了一个示例：假设函数 h(x) = f(g(x))，则 h'(x) = f'(g(x)) * g'(x)（此处的 f'、h' 和 g' 分别是 f、h 和 g 的导数）。有了这条规则，组合函数就不需要任何麻烦的计算，只要了解如何获取每个组成部分的导数即可。

举个例子，假设你有 sin(x**2)，则只需知道组成函数 x**2（即 2*x）和 sin(x)（即cos(x)）的导数，即可求出结果为：cos(x**2) * 2x。

事实证明，链式法则对于大型表达式图求导非常有用。你不需要仔细研究如何对庞大和过于复杂的函数求导。你只需要已理解的构建块，并且它们是组合而成的。

下面，我们将链式法则应用于表达式图。

将链式法则应用于图

首先，我们从一个 Value 节点开始。对于给定的节点，我们可以执行链式法则的一步（伪代码）：

此处 wrt 的意思是“对于”。求每个子节点对于某个子节点的导数，这一点非常重要。

我们不只是设置 child.grad，而是使用了 +=，原因有两个：

一个子节点可能被多个父节点使用，在这种情况下，子节点会影响到所有父节点。

批处理，暂时无需在意。

为了更具体地说明，下面我们来看看 Karpathy 实现 * 的导数的方法。在数学计算中，如果 f(x,y) = x*y，则 f'(x, y) = 1*y（对于 x）且 f'(x, y) = x*1 （对于 y）。代码如下：

这意味着，对于每个子节点而言，我们将使用另一个子节点的数据并（由于链式法则）将其与父表达式的梯度相乘。也就是说，self.grad（左侧）是使用 other.data（右侧）调整的，反之亦然。我们通过上述步骤成功地将数学计算转换成了代码。求导数，应用链式法则，然后加到子节点的梯度上。

到这里，我们建立了一个函数来求操作节点的导数，但我们需要对整个图求导。

遍历图并不像遍历树那么简单。你需要避免重复访问同一个节点，并保证在父节点之前访问子节点（正向模式）或在子节点之前访问父节点（反向模式）。难点在于，虽然我们不会重复访问同一个节点，但访问会更新该节点的子节点（而不是节点本身），并且多个节点可能共享子节点，因此子节点的梯度可能会被多次更新。这都是正常情况。

为此，我们必须引入拓扑排序。

拓扑排序和图转换

图的拓扑排序能保证图中的子节点永远先于父节点被访问。一般来说，只有当图中没有环路时才能使用拓扑排序，幸运的是，我们前面已经假设图没有环路。

下面是 Value 图的拓扑排序示例。为了简洁起见，我们使用了嵌套函数 build_topo，但这并不是绝对必要的。

为了说明上述代码的工作原理，我们可以针对一个非常简单的表达式图 1*2 进行拓扑排序。

在这段拓扑排序中，为了计算值 3，首先我们必须计算值 1 和 2。值 1 和 2 的计算顺序并不重要，但二者都必须在 3 之前计算。

以上，我们确定了图的遍历顺序，下面可以着手反向传播了。

将拓扑排序应用于反向传播

我们可以利用上述介绍的链式法则和拓扑排序，在图上进行反向传播。下面是 micrograd 的实现代码。首先构建一个拓扑排序，然后对其进行反向操作，将链式法则应用于每个 Value。

通常，我们会在损失函数的结果 Value 上调用 Value.backward。

你可能在想为什么我们在进行反向传播之前将 self.grad 设置为 1，你可以仔细想一想。

整合

此处我不打算深入细节，只是粗略地介绍一下如何通过简单的训练，使用基于 MLP 的分类器解决 MNIST 数字识别问题。下面的代码不能直接运行，其中缺少图像加载支持代码和损失函数。超参数（批量大小等）是任意设定的，且未经调整。完整的训练代码和添加 exp/log/Max 的引擎改动，请参见 GitHub 代码库。

训练代码：https://github.com/tekknolagi/micrograd/blob/534ab3c884e66c8a325e0a8f3ed278656a616002/mnist.py

相应的引擎改动：https://github.com/tekknolagi/micrograd/blob/534ab3c884e66c8a325e0a8f3ed278656a616002/micrograd/engine.py

在上述代码片段中，MLP (model = MLP(...)) 会在各层中构建一堆神经元，并将一些权重初始化为 Value 的实例，但它还没有构建图。只有被调用时（model(image.pixels)），它才会构建图并执行所有点积。然后，在计算损失时，我们在此基础上构建更多的图。这就是前向传播。

接着，我们反向传播，利用损失调用 backward()。

然后，我们通过梯度来调整所有权重。

最后，请不要忘记将梯度初始化为零！

性能问题

用 CPython 运行这段代码会非常慢，感觉上计算一张图的前向传播大约需要一秒钟，我们还需要反向传播，而且我们必须对 6万 张图进行几个 epoch 的处理。最终花费的时间会过长。

我们可以按照大家的建议，尝试使用 PyPy。这样每秒可以处理几张图，但仍然不够快。

顺便说一下，我们的旧项目 Skybison 比 CPython 和 PyPy 都快得多！经过分析后，我们发现性能的主要痛点是函数创建（在 Skybison中有点慢），但如果将内部函数 _backward 放到顶层，问题就会消失。因此，很明显拓扑排序的集合查找是配置文件中最慢的部分。之后是所有临时 Value 对象的垃圾回收。

另外，将内部函数放到顶层也可以极大地提高 PyPy 的速度，并且比 Skybison 更快。

我认为所有运行时的痛点是：

每次前向传递都会重新创建图，因为所有的 Value及其 _backward 函数必须重新分配。

Neuron.__call__ 中的 zip 也存在大量内存分配和迭代开销。

由于指针追逐、函数调用以及集合/列表内存分配和操作，每次反向传播都会进行拓扑排序。

正常的 Python解释器开销。

但根据多年的经验，我认为首先应该实际测量一下，而不是在黑暗中盲目优化。

使用分析器检查

EmeryBerger 和他的团队发布了一款出色的 Python 分析工具，名为 Scalene。使用时，你可以直接运行 scalene yourprogram.py（代替python3yourprogram.py）。程序运行完成后（或者按Control-C键），将弹出一个本地托管的小型网站，其中包含分析信息。

我在 micrograd MNIST 上运行了 Scalene，结果如下：

图：Scalene 分析器输出的 micrograd 分析结果屏幕截图。

我们可以看到许多 Value 的内存分配，而且 self._prev 是一个集合，甚至有可能造成内存泄漏。特别是，我们还可以看到很多 + 和 * 操作，因为 __add__和 __mul__ 分配了很多内存。

看看内存使用列，曲线呈向上向右的趋势，这不是我们想要的。似乎为每个 Value 创建 _prev 元素的set 的过程占据了大量时间。

如果你是守旧派，不信任新的分析工具，那么甚至可以使用 perf 确认这些观察结果。你可能需要根据 Python 发行版安装调试符号（我使用的是 Ubuntu 的 python3.10-dbg），然后运行 perf record python3 yourprogram.py。我得到的结果如下（省略了 0.5% 以下的记录）：

gc_collect_main占总体时间的 37% 是一个巨大的危险信号。其次，下面的其他函数（deduce_unreachable和所有的 _traverse 函数）看起来也与垃圾回收相关，这意味着程序占用了太多内存。所以，Scalene 和 perf 的分析结果似乎是一致的。

如果去掉 set(_children)，仅使用元组（这似乎不会影响正确性），则时间占用会相对分散一些。

还有一个简单的方法是将 __slots__ 添加到 Value 类。属性字典是我能想到的唯一分配字典的地方，所以也许我们可以解决这个问题。果然添加 __slots__ 后，dict_traverse 就消失了。

最后，我们还可以尝试删除嵌套函数分配，这样就可以消除 func_traverse。不过这项优化要比前两个更加繁琐一点。

这些小改动不会改变程序的整体架构，所以也不会带来大量的数学运算和图遍历工作。

那么，我们应该怎么办呢？

解决方案

提高程序运行速度的最佳方法是减少工作量。垃圾回收太多？则减少内存分配。递归太多？则减少拓扑分类。开销太大？则减少解释的工作量。更详细地说，我提出的解决方案是：

重用输入之间的图结构。避免每次都重新构建 Value 图，复制新输入并执行前向传播和反向传播。

由于不修改图，因此也无需重新拓扑排序。顺序保持不变，有利于前向传播和反向传播。

归根结底，Value抽象并不重要。如果我们知道遍历的顺序并使用 IEEE-754 双精度，就应该将拓扑排序及其操作编译为 C 或更简单的东西。

这与我们了解的编译器知识相一致：如果可以在程序允许的语义中冻结一些动态，就可以提升性能。由于图是静态的，所以我们可以采用这种做法。

编写编译器

这个编译器的目标是，为 micrograd 编写一款非常小且非常适合的编译器，不需要重新设计架构。

我们可以编写一种字节码编译器，并通过这个编译器去除所有函数调用以及重复的树遍历和指针追逐。这样就能提升性能。但不幸的是，我们仍然有一个解释器循环，并且该解释器是用 Python 编写的，这会产生大量开销。

为此，我们将进一步将这些代码编译为 C。最终目标是编写一个 Python C 插件，我们可以导入并使用它来代替 micrograd 的解释版本。

该项目的最初版本是将 MLP、Layer 和 Neuron 类直接编译为 C，但不幸的是，它的可扩展性不是很好：修改模型的架构需要编写新的编译器。此外， 它也不支持反向传播，只是对推理有帮助。

因此，我们需要为 Value 图编写编译器。这意味着，只要机器学习架构使用 Values，那么任何人都可以毫不费力地使用这款编译器。你只需要为它写一个解释器。

前向传播

由于我们使用了拓扑排序，所以不妨在前向传播和反向传播中使用它。那么，我们只需要编写一个一次只处理一个 Value 的编译器。写法如下：

（假设 data 是我们稍后将创建的一个大小适当的双精度数组。）

上面的代码把图变成了线性。这有点像我们之前看到的拓扑排序，但是使用 C 代码编写的。这种策略之所以有效，是因为我们没有循环，也没有重新定义 Values。每个值设置一次，而且这段代码即使包含内存加载和存储，也应该比 Python 中的指针追踪和函数调用快得多。

我们可以编写一个类似的解释版本，每种操作都有自己的方法（__add__、__mul__ 等），但将编译器全部呈现在一个方法中更容易。因此我添加了一个编译函数。下面的示例实现了常量值（op==''）和加法（op=='+'）：

其他运算符的写法也一样。例如，你可以试试看如何实现 ** 或 exp。请注意，** 需要存储额外的数据或某种特殊的处理。

你可能会注意到，这种编译策略需要为 Values 分配 ID。为此，我在__init__ 函数中添加了一个 _id 字段，它是__init__ 函数中的自动递增计数器。具体的实现并不重要，你只需要知道每个 Value 对象都有一个唯一的 _id 即可。

我的编译器实现所有的操作大约为 40 行代码，甚至包括一些小的即时优化。但这个编译器是前向传播。反向传播呢？我们也需要加快训练速度。向后传播一定更为复杂，是吗？

反向传播

实际上，反向传播的复杂度与前向传播差不多。我们只需要逐行修改反向传播函数（所有的 _backward 实现）。

例如，我们需要修改 * 的反向传播。我添加了一些辅助函数，这样代码行数更少，而且看起来更像解释版本。与前向传播一样，所有运算符都在一个方法：backward_compile。

（与前向传播一样，我们假设 grad 是稍后即将创建的一个大小适中的双精度数组。）

下面，我们来看看如何应用：

很奇怪，为什么 x (grad[6]) 和 y (grad[7]) 没有反向传播代码？因为它们没有子节点，而本身由父节点 z (grad[8]) 调整。我前面曾提到过，访问节点会调整该节点的子节点。

我的编译器实现反向传播大约为 30 行代码，甚至比前向传播还短，非常整洁。

如此，我们就完成了编译器的编写。恭喜！本文最复杂的部分已经结束了。其余都是一些小细节和 Python C-API 的具体实现。

更新权重

在实现了反向传播后，我们需要通过它们的梯度来调整权重。此处的代码只需将 Python 代码机械地翻译成 C。为了方便比较，下面是解释版本：

它会在运行期间遍历并调整模型参数。相比之下，编译版本在编译时进行迭代，而在运行时仅执行减法操作：

如果不考虑 assert，上述代码的长度与 Python 相差无几。

设置输入

当输入不是整数和浮点数等简单数据类型时，将 Python 代码输入到 C++ 会有点棘手。理想情况下，我们生成的机器学习代码能够与 Python 共享内存，这样就可以避免来回复制数据，但这种实现并不简单，因此我们需要采用略麻烦的做法。

我们来添加一个函数 set_input，将黑白像素数据放入字节数组中，并将每个像素复制到 data 数组的每个槽中。虽然这种方法相对较慢，但肯定不会成为管道中的瓶颈。

在这个例子中，inp 是输入数组。与 micrograd 的解释版本不同，我们不会在每次迭代中创建新的输入 Values。这意味着，我们必须预先分配机器学习模型输入和输出的 ID 范围：

请注意，由于 inp 和 exp 是任意选择的，所以每个 Value 节点的 data 或grad 字段都包含垃圾数据。但是，生成的 C 代码并不会使用这些 Python 值。我们关心的是 _op 和 _prev 字段表示的图结构。

为了使用 Python 中的 C 代码，我们必须使用 C-API 来创建 Python C 插件。

PythonC 插件

通过一堆代码更新 data 和 grad 度数组很有趣，而且它是一个完整的编译器，但目前还不能发挥作用。我们需要将这些代码包装到函数中（我为它们取名为 forward、backward、update和 set_input），并允许 Python 驱动程序访问这些函数。我们不想完全使用 C！

大部分的代码很简单（只需要添加 print("void forward() {" 等代码），但部分代码需要用到 Python 的内部知识。

例如，下面是包装 forward 函数的代码：

这是一个fastcall C-API 函数的示例，这意味着它会通过一个数组获取参数。我们必须像下面这样注册这个函数：

接下来，我们创建一个可供 Python 导入的模块描述，这样就可以在导入时创建模块对象：

然后，我们来创建 PyInit_nn 函数。如果 Python 的原生导入器在.so 中找到模块，并且这个模块拥有 PyInit_XYZ函数，则调用它来创建模块对象。

到此，我们的编译器就基本编写完成了。接下来的主要工作是模型的训练和推理。

正确吗？速度提升了吗？

这是两个不同的问题，如果代码产生错误的输出，那么性能就没有任何意义了。

正确性

测试编译器有些棘手。编译器的很多部分不仅必须能够单独工作，同时必须与其他部分协同工作。值得庆幸的是，在这个示例中，我们的编译器非常小，只包含少量基本操作。因此，为生成的 C 代码编写单元测试并不太难。

此外，我们还可以针对同一段代码的解释版本和编译版本输出的数字进行一些并行测试。如果二者都有一定的误差范围，则我们可以认为编译器是正确的。不过，我不建议使用 MNIST。解释版本太慢，而单元测试应该能够很快地运行。或许可以试试看XOR。

值得庆幸的是，CPython 的 float 使用了宿主系统的浮点实现，因此我们无需额外的努力即可获得与 C 相同的数字行为。

性能

在我的机器上，训练从每秒 1 个图像（解释版本）上升到了每秒 > 1000 个图像（编译版本），性能提升大约为 1 千倍！不过，编译版本会产生一些前期的开销，因为你必须编译 C 代码。如果使用 TCC（一种速度非常快的 C 编译器），那么可以获得非常好的性能。我的编译时间大约为半秒，每个 epoch 大约需要 45 秒。如果使用 Clang（一种相对很慢的 C 编译器），则可以获得更好的性能。具体的数字如下表所示：

不管怎样看，这都是一个巨大的胜利。我认为我们成功了！

完整的编译器代码、编译器包装器和训练代码，请参见 GitHub：

编译器代码：https://github.com/tekknolagi/micrograd/blob/c15b6b8fd373c48014be369c4f7bd0917932a53b/micrograd/engine.py

编译器包装器和训练代码：https://github.com/tekknolagi/micrograd/blob/c15b6b8fd373c48014be369c4f7bd0917932a53b/test.py

总结

神经网络由静态数据流图表示，可向前或向后执行。这意味着，它们有点像遍历树的解释器。这也意味着，将树编译为较低级别的表示可以加快程序的运行速度。