1. 为什么从零开始手搓一个计算图1.1 计算图到底是什么计算图本质上是一个有向无环图。图的每一个节点代表一次计算或者一个数值载体边代表数据依赖关系。举个最简单的例子a b * c这条表达式展开就是三节点的结构乘法节点依赖b和c加法节点依赖a和乘法节点的输出。整个前向传播就是按照依赖关系把每个节点依次算出来而反向传播则是沿着相反的方向用链式法则把梯度一层层传回去。我第一次接触这个概念的时候有个误区以为计算图只有神经网络才用得上。后来才意识到任何输入-计算-输出的流程都能用这张图描述。自动微分、符号求导、程序改写干的都是类似的事。只不过在深度学习场景里我们需要对一个包含大量参数和中间结果的复杂表达式反复求梯度计算图这个结构几乎是必然解。主流的训练框架内部做的事情跟你现在要手搓的东西在本质上没有区别只是工程上做得更极致、更高效而已。1.2 手搓一个图能获得什么直接使用框架当然很快但框架把太多细节藏起来了。你调用一个损失函数它自动就帮你把梯度算好了可一旦遇到梯度异常、自定义算子、或者你想搞清楚某个数值问题到底出在哪一层的时候不了解底层机制就会非常被动。我自己最深的体会是当你能独立写出一张可以自动求导的计算图你对反向传播这四个字的理解会完全不一样。以前是背公式知道链式法则长什么样写完之后你会建立起一种直觉——每个算子的梯度应该如何回传、多个分支的梯度为什么要相加、哪些操作在求导时会改变数据的形状。这种直觉是看多少文档都补不回来的。而且这张图是你自己的想加什么算子就加什么算子想打印哪一步的中间结果就打印哪一步自由度比用现成框架高得多。这一篇的目标很具体用 C 实现计算图的核心引擎包括节点定义、前后向传播、拓扑排序和基础算子然后用 Python 封装出一套用户友好的接口最终跑通一个线性回归的训练例子并且用数值梯度验证结果。整个过程我会把代码一步步拆开讲包括踩过的坑。2. 整体设计先把这张图的骨架想清楚2.1 数据结构选型写计算图第一步不是写代码而是想清楚节点里到底要存什么。我的设计里一个节点至少包含以下几样东西节点类型常量、变量、还是某个算子输入节点列表这个节点的计算依赖谁输出节点列表哪些节点的计算依赖它前向计算值节点在当前输入下的结果梯度值反向传播时累积到节点上的梯度为什么输出节点列表也要单独存一份因为反向传播时要从输出往输入回传如果每次都遍历全图找谁依赖了我复杂度会很难看。提前把邻接关系双向记好前向和反向就都能顺着指针走效率高代码也直白。生命周期管理上我直接用std::shared_ptr。每个节点可能被多个上层节点引用图本身又可能反复扩展引用计数能自动解决大部分释放问题后面调试期省了很多头疼。代价是循环引用要小心但我设计的图是严格有向无环的节点只会指向依赖方向不会出现互指成环的情况。2.2 前向传播怎么做前向传播的核心难题是执行顺序一个节点必须在它的所有输入都算完之后才能算。最简单可靠的做法是拓扑排序。我每次构建完图会对所有节点做一次 Kahn 算法不断把入度为零的节点挑出来执行得到一个有序的执行序列。之后每次前向直接按这个序列遍历执行即可。这里有一个工程上的优化思考很多框架在前向的时候并不显式做拓扑排序而是靠节点被调用时递归计算输入的方式。这种方式写起来更简单但会有两个问题一是同一节点被多个上层节点使用时会被重复计算冗余严重二是递归深度太深的时候栈容易爆。所以我选择显式拓扑排序加顺序遍历不仅稳定后续做算子融合之类的优化也更顺手。每个算子的前向实现我把它封装成统一的函数签名输入一组节点的值输出这个节点的值。这样图引擎本身完全不关心算子内部逻辑只管调度。2.3 反向传播的核心思路反向传播的数学基础是链式法则假设有链式关系y f(u)u g(x)那么dy/dx (dy/du) * (du/dx)。在计算图里一个节点可能同时被多个上层节点使用例如c a b同时d a * 2那么a的梯度来自c和d两个方向的贡献必须相加。这就是为什么反向传播里反复强调梯度累积。基于这个原理每个算子只需要实现自己局部的反向逻辑拿到从上游传来的梯度计算出对每个输入的梯度贡献再累加到输入的grad字段里。图引擎负责的事情很简单——按照拓扑排序的逆序遍历所有节点逐个调用算子的反向函数。至于复杂的链式组合让图自己去逐层传递就好每个算子只需要管好自己这一跳。这个设计的优雅之处在于新增一个算子只需要写两个函数前向和反向图的调度逻辑完全不用改。下面会详细说 C 代码里怎么落地。3. C 核心实现引擎部分3.1 节点与图的基本代码我先把基础类型定义出来。为了控制篇幅这里用一个简化版的Tensor本质上就是二维vector实际项目里你可以换用任何你顺手的矩阵库接口保持一致即可。#include vector #include memory #include functional #include unordered_map #include algorithm using Matrix std::vectorstd::vectordouble; struct Node; using NodePtr std::shared_ptrNode; enum class OpType { CONSTANT, VARIABLE, ADD, SUB, MUL, DIV, MATMUL, RELU }; struct Node : std::enable_shared_from_thisNode { OpType op; std::vectorNodePtr inputs; std::vectorNodePtr outputs; Matrix value; Matrix grad; int id; Node(OpType t) : op(t) {} };为什么用enable_shared_from_this因为构建节点时需要让节点自己拿到shared_ptr形式的指针塞进别人的outputs列表里。如果不这么做就得在外部统一管理节点创建和连接接口会繁琐。这个写法在二叉树的增删里我也常用属于引用计数体系下的常规工具。接下来是图类负责节点的注册、拓扑排序、以及前后向执行class Graph { public: NodePtr addNode(OpType op, const std::vectorNodePtr inputs) { auto node std::make_sharedNode(op); node-id nextId; node-inputs inputs; for (auto inp : inputs) { inp-outputs.push_back(node); } nodes.push_back(node); return node; } NodePtr addConst(const Matrix val); NodePtr addVariable(const Matrix val); NodePtr addMatMul(NodePtr a, NodePtr b); NodePtr addRelu(NodePtr a); void forward(); void backward(); private: std::vectorNodePtr nodes; int nextId 0; std::vectorNodePtr topologicalOrder(); };这个图类的设计基调是图自己知道所有节点。这也带来一个注意点如果构建了一个节点但不挂进图里直接手动连接那它永远不会被拓扑排序发现前后向都不会执行它。我在早期版本里就踩过这个坑那时候允许节点从外部创建结果经常出现明明加了节点输出却不对的情况。后来强制所有节点必须通过Graph::addNode创建问题就绝迹了。3.2 算子的前向与反向实现统一调度是核心思想。我把算子实现注册到一张表里前向和反向都通过函数指针调用。这样引擎代码只关心节点类型是什么就调用什么函数而不需要写一长串if-else。using ForwardFn std::functionMatrix(const std::vectorMatrix); using BackwardFn std::functionstd::vectorMatrix( const std::vectorMatrix inputValues, const Matrix gradOutput); struct OpInfo { ForwardFn forward; BackwardFn backward; };这里反向函数的输入为什么不仅要gradOutput还要inputValues因为某些算子的梯度计算依赖前向输入的值。最典型的是乘法y a * bdy/da bdy/db a。如果反向函数拿不到前向的a和b梯度就算不出来。这个细节在设计算子接口时必须从一开始就留好位置事后补接口会牵连很多调用方。具体算子的实现视情况而定。加法算子的前向就是把输入逐元素相加反向把上游梯度原样分发给每个输入。矩阵乘法则稍微复杂矩阵乘法的局部梯度公式是若Y A * B则dL/dA dL/dY * B^TdL/dB A^T * dL/dY。这里面最容易犯的错是维度搞反。我写第一版矩阵乘反向时不自量力地凭记忆写了个交换顺序的乘法结果跑出来的梯度形状完全不对。后来老老实实把维度标注清楚每次写完先拿小矩阵手算一遍再跑代码。这种基础算子的正确性直接决定上层所有结果宁可在这一步慢一点。ReLU 算子的逻辑也很典型前向就是max(0, x)反向则需要记住前向时哪些位置是负数负数的位置梯度直接置零正数的位置梯度原样通过。这也是反向依赖前向输入的另一个例子。3.3 拓扑排序与执行拓扑排序的代码我直接基于 Kahn 算法std::vectorNodePtr Graph::topologicalOrder() { std::unordered_mapNodePtr, int inDegree; for (auto node : nodes) { inDegree[node] node-inputs.size(); } std::vectorNodePtr result; std::vectorNodePtr queue; for (auto kv : inDegree) { if (kv.second 0) queue.push_back(kv.first); } while (!queue.empty()) { NodePtr cur queue.back(); queue.pop_back(); result.push_back(cur); for (auto out : cur-outputs) { if (--inDegree[out] 0) { queue.push_back(out); } } } return result; }这里有一个小的经验入度为零的节点用栈还是队列存储只影响同层节点之间的执行顺序不影响最终的正确性。但如果你是做数值稳定性调试希望同一个图的执行顺序每次都完全一样那就要使用稳定的容器而不是无序的哈希表遍历。我在unordered_map上吃过亏同样的代码两次运行执行顺序不同某些浮点累加的结果就有极微小的差异。排查到最后才发现是哈希表的迭代顺序不稳定。所以这个函数里我用了节点 id 排序或者直接用vector保存入度信息保证顺序可复现。前向执行很简单按拓扑顺序依次调用算子的forward函数把结果写进node-value。反向则是逆序遍历每访问一个节点就把它的grad当作上游梯度调用算子的backward得到对每个输入的梯度贡献累加到输入的grad上。后端这些思路确认无误后上层就可以开始封装 Python 接口了。4. Python 封装从核心引擎到顺手接口4.1 直接做绑定还是重新实现写 Python 接口有两条路线一条是给 C 核心写 Python 绑定另一条是用纯 Python 照着同样思路实现一遍。我这个项目两者都试过。最终保留的是纯 Python 实现原因很实际当你的核心逻辑已经吃透之后纯 Python 版的代码可读性更高调试也方便用户的每个算子都能直接钻进去看源代码。C 绑定主要在性能敏感的场景才有必要。所以下面的 Python 实现结构上完全模拟 C 核心的设计思想但语法上利用了 Python 的灵活特性。这个项目的定位是理解原理而不是追求极致性能纯 Python 反而更合适。4.2 Tensor 类与运算符重载用户的体验应该尽可能像操作普通数字一样。Python 的运算符重载在这方面非常顺手class Tensor: def __init__(self, data, requires_gradFalse): self.data np.array(data, dtypenp.float64) self.requires_grad requires_grad self.grad np.zeros_like(self.data) if requires_grad else None self.creator None # 产生这个 Tensor 的算子 self.children [] # 依赖这个 Tensor 的后续节点 def __add__(self, other): if not isinstance(other, Tensor): other Tensor(other) return add(self, other) def __mul__(self, other): if not isinstance(other, Tensor): other Tensor(other) return mul(self, other) def __matmul__(self, other): return matmul(self, other) def backward(self): # 先做拓扑排序再逆序回传 order topo_sort(self) for node in reversed(order): node._backward_fn()creator记录的是谁生成了我children记录的是我生成了谁。这跟 C 版本里的inputs和outputs一一对应只是换成了 Python 的命名习惯。这里有一个值得思考的设计点为什么在 Python 版本里不显式建 Graph 对象而是让每个 Tensor 自己携带创建关系答案是 Python 的 GC 和引用跟踪更适合这种需要从输出回溯到所有祖先的模式。你在做loss.backward()时并不需要知道全图有多少节点只需要从loss这个节点出发向前递归就能遍历到所有依赖路径。C 版本显式建图是因为它需要精确控制生命周期和顺序Python 版本则把图的构建融入计算过程用起来更像原生代码。4.3 算子的构建方式每个算子都要同时注册前向函数和反向函数。我采用了一个简单的注册表设计def add(a, b): result Tensor(a.data b.data, requires_grada.requires_grad or b.requires_grad) if result.requires_grad: result.creator (add, a, b) a.children.append(result) b.children.append(result) return result反向的时候根据creator的类型做不同的梯度分配def _backward_add(node): _, a, b node.creator if a.requires_grad: a.grad node.grad if b.requires_grad: b.grad node.grad注意这里是而不是。这就是我在 C 部分提到的梯度累积规则。如果写成赋值当一个输入被多个下游节点共用时后执行的节点会覆盖先执行的梯度结果就错了。我最早写这个逻辑时用的是测试一个简单的共享输入表达式(a b) * (a c)时a的梯度怎么都不对。后来打印每一项的梯度贡献才发现两个分支的贡献只有一个被保留了。改成累加后立刻就好了。这个坑几乎每个人写自动求导都会踩一次。5. 跑通完整例子从数值验证到线性回归5.1 数值梯度验证手写自动求导正确性验证是第一位的。我每次给图新增算子之后都会用数值差分法做一次交叉验证。数值梯度的原理很简单对某个变量做微小的扰动用(f(xε) - f(x-ε)) / (2ε)近似导数然后跟计算图算出来的梯度对比。def numerical_gradient(f, x, eps1e-6): grad np.zeros_like(x, dtypenp.float64) it np.nditer(x, flags[multi_index]) while not it.finished: idx it.multi_index old x[idx] x[idx] old eps plus f(x) x[idx] old - eps minus f(x) x[idx] old grad[idx] (plus - minus) / (2 * eps) it.iternext() return grad这个验证方法非常土但极其有效。我建议每一个自定义算子都过一遍这个测试尤其是涉及维度变换的算子比如矩阵乘法、卷积、reshape因为形状错乱的时候数值梯度能一针见血地告诉你梯度值完全不对。我个人测过一个很经典的表达式loss sum((a b c) ** 2)其中a、b、c都是不同形状的变量。数值梯度和反向传播梯度的相对误差在1e-8以下才算通过。如果你跑出来误差在1e-5量级往往不是算法错而是epsilon取值不好或者浮点精度问题如果误差直接是个位数甚至几百倍那就是实现错得离谱优先检查维度匹配。5.2 线性回归训练实现通过验证后我用一个线性回归做端到端的测试。虽然简单但它能覆盖计算图的完整生命周期前向预测、计算损失、反传梯度、更新参数。import numpy as np # 模拟数据 y 3.0 * x 1.5 x_data np.random.randn(128, 1) y_true 3.0 * x_data 1.5 w Tensor(np.random.randn(1, 1) * 0.1, requires_gradTrue) b Tensor(np.zeros((1, 1)), requires_gradTrue) learning_rate 0.1 for step in range(100): pred x_data w b diff pred - y_true loss (diff * diff).mean() loss.backward() w.data - learning_rate * w.grad b.data - learning_rate * b.grad w.grad np.zeros_like(w.grad) b.grad np.zeros_like(b.grad) if step % 20 0: print(fstep {step}, loss {loss.data:.6f}) print(ftrained w {w.data.ravel()[0]:.4f}, b {b.data.ravel()[0]:.4f})这个例子跑通之后说明整条链路是通的。你可能会疑惑为什么每次都手动把grad清零因为我们的 Tensor 定义里梯度是累加的下一轮迭代必须手动清零。主流框架通常会在backward里自动处理或者提供zero_grad方法。这个细节反映了设计选择自动清零方便但容易掩盖问题手动清零明确但容易漏。在早期学习阶段手动清零能强迫你理解梯度的生命周期。6. 常见问题与排查实录6.1 梯度全是 NaN这是自动求导里最常见的故障原因通常有两个一是前向计算里出现了除零或者log(0)反向时1/x或1/log(x)这类梯度直接变成无穷大二是梯度爆炸多层累乘之后数值溢出。排查思路我总结了一个固定流程先看前向输出是否异常再看反向第一步的梯度然后逐层往前打印梯度值定位到哪一层开始出现inf或nan。不要上来就怀疑图引擎本身大部分情况是某个输入数值范围不对。比如初始化权重太大经过若干层矩阵乘后数值指数增长梯度自然爆炸。把权重初始化标准差调小到0.01量级多半就能解决。6.2 反向传播时某些节点梯度为 0节点梯度为零不一定是 bug。如果一个变量对最终输出确实没有贡献——比如loss a * 0 b里的a——梯度为零是正确的。但如果你确认它应该对输出有影响梯度却为零就要检查梯度累积是否被覆盖、拓扑排序是否包含了这个节点、以及算子反向函数是否把梯度传到了正确的输入位置。我调试过一个具体案例自定义的某个加权求和算子反向时不小心把第一个输入的梯度写到了第二个输入的grad上。因为两个输入形状恰好一样输出又恰好与第二个输入有关整个结果看起来非常合理直到换了测试数据才发现不对。所以调试时一定要把输入-梯度的对应关系列成一张表逐一核对。6.3 numpy 操作容易出现的坑Python 实现里我大量用了 numpy有几个点必须时刻注意确保所有 Tensor 内部数据是float64整数矩阵做梯度运算时会直接报错或者静默截断。反向时修改grad要小心 broadcasting。numpy 的广播机制在形状不完全相同时会自动扩展这会导致梯度分配到错误的位置。比如形状(3, 1)和(3,)相加反向时对后者求梯度必须sum掉广播引入的维度。不要让Tensor的data和某个外部数组共享内存。一旦误用外部数据的意外修改会污染计算图的状态。6.4 调试笔记和技巧我给自己的项目强制定了一条规矩每个算子旁边都保留一个数值验证样例用一个小到能心算的矩阵跑一遍把期望梯度直接写在注释里。这样任何人包括三个月后的我自己改代码时跑一遍样例就知道有没有破坏算子逻辑。这个习惯后来帮了大忙有一次想优化矩阵乘法的反向实现直接跑样例就发现新版代码在某个边界维度上算错了。另外建议做一个打印梯度的辅助函数递归打印某个节点所有祖先节点的值、形状和梯度按拓扑序输出。调试的时候比起到处打print这个工具能让你一眼看到整个图的状态分布。写在最后这次把计算图从零手搓了一遍之后我再回去看那些带自动求导的框架心态完全不同了。以前觉得反向传播是一个黑盒魔法现在知道它就是一张 DAG 上的逆序遍历加上一堆按链式法则实现的局部求导函数。这个项目做完后我给自己留了三个后续方向一是加入更多算子比如卷积、池化、softmax二是实现更高效的内存复用避免每个节点都保存完整的中间矩阵三是把一个简单的优化器比如 Adam也纳入图的管理让参数更新本身成为图的一部分。如果你也在手搓这个过程我的建议是不要追求一次写完所有算子先把加法和乘法跑通再用数值梯度死死守住正确性然后逐步扩展。这个小项目的价值不在代码量而在你亲手构建起的那份对自动求导的直觉。
