1. 为什么我要从零写一个深度学习框架动手写框架这件事我在刚入行那几年是想都不敢想的。那时候觉得 TensorFlow、PyTorch 这种体量的东西背后是几百个博士和几千个工程师的集体智慧跟我一个普通开发者没什么关系。直到有一次我在调试一个自定义算子的时候发现底层报错信息完全看不懂翻源码翻了整整两天才定位到一个张量步长的问题。那次经历让我意识到如果一直停留在“调包”的层面遇到真正棘手的问题时就会非常被动。于是我开始尝试自己写一个极简的深度学习框架。目的不是造一个能替代 PyTorch 的轮子而是通过亲手实现每一个核心模块把自动求导、计算图、张量运算这些概念从“知道”变成“理解”。这个系列的第一篇我想先把架构设计这件事聊透。因为架构决定了后面所有代码的组织方式如果一开始方向就偏了写到后面会非常痛苦。这篇文章适合谁看如果你已经会用 PyTorch 或 TensorFlow 训练模型但对反向传播的底层实现只有模糊印象如果你能看懂 Python 代码但不确定一个框架从零开始应该先写什么、后写什么如果你也想通过造轮子来加深理解但不知道从哪里下手——那这篇内容应该能帮到你。我会把架构设计的思路、模块划分的理由、以及我在这个过程中踩过的坑尽量完整地分享出来。2. 整体架构设计思路拆解2.1 先想清楚框架到底要解决什么问题在动手画架构图之前我强迫自己先用一句话回答这个框架的输入是什么输出是什么中间经历了什么。答案其实很朴素——输入是数据和模型定义输出是梯度中间的核心是“如何自动地、正确地计算出梯度”。围绕这个核心框架需要具备几个基本能力。第一是张量运算因为所有神经网络的计算最终都归结为矩阵乘法、加法、激活函数这些操作。第二是自动求导这是深度学习框架区别于普通数值计算库的关键。第三是计算图的构建与执行它决定了运算的顺序和依赖关系。第四是参数管理与优化器负责根据梯度更新权重。把这四件事想清楚之后架构的轮廓就出来了。我选择的是基于动态图的 define-by-run 模式也就是每次前向传播时实时构建计算图而不是像早期 TensorFlow 那样先定义静态图再执行。这个选择的原因后面会详细说但简单来讲动态图更符合 Python 的使用直觉调试起来也方便得多。2.2 为什么选择动态图而不是静态图静态图的思路是“先定义、后运行”。你先把整个计算流程描述成一个图结构框架再对这个图进行优化和执行。这种方式的优势在于图一旦确定就可以做大量的编译期优化比如算子融合、内存复用、并行调度。但它的缺点也很明显调试困难因为你看不到中间结果灵活性差因为图的结构在运行前就固定了。动态图则是“边运行、边建图”。每次前向传播的时候框架根据实际执行的操作动态地记录计算历史。这样做的好处是你可以像写普通 Python 代码一样写模型用 print 打印中间结果用断点调试非常直观。代价是失去了部分全局优化的机会但对于一个学习性质的框架来说可调试性和可理解性远比极致性能重要。我最终选择动态图还有一个很实际的原因实现起来代码量更少逻辑更清晰。静态图需要你先设计一套图描述语言或者 IR然后再写一个执行引擎复杂度直接翻倍。而动态图只需要在每次运算时记录一下操作和输入输出关系后面反向传播时沿着记录链回溯就行了。2.3 核心模块的划分与依赖关系整个框架我划分成了五个核心模块它们之间的依赖关系是自底向上的。最底层是张量模块负责数据的存储和基础运算。它不依赖任何其他模块是整个框架的基石。往上是自动求导模块它在张量运算的基础上记录计算历史并实现反向传播。再往上是计算图模块负责管理节点之间的依赖关系实际上自动求导模块和计算图模块是紧密耦合的我倾向于把它们放在一起考虑。然后是神经网络层模块它基于张量和自动求导封装了全连接层、激活层、损失函数等常用组件。最上层是优化器模块它接收自动求导计算出的梯度按照一定的策略更新参数。这种分层设计的好处是每一层只依赖它下面的层不会出现循环依赖。我在早期版本中曾经把优化器和自动求导混在一起写结果代码很快就变得难以维护。后来把优化器独立出来之后整个结构清晰了很多。2.4 一个关键决策张量是否要支持 GPU这个问题我想了很久。支持 GPU 意味着需要写 CUDA 核函数或者调用 cuBLAS 之类的库代码复杂度会急剧上升。而且对于学习框架设计原理来说GPU 并不是必须的。CPU 上的实现同样能讲清楚自动求导和计算图的逻辑。所以我最终决定第一版只支持 CPU用 NumPy 作为底层的数值计算库。这样做的好处是我可以把精力集中在框架本身的架构上而不是被 GPU 的内存管理和核函数编写分散注意力。当然架构上我会预留扩展点比如在张量类中抽象出一个 device 属性后面如果想加 GPU 支持只需要实现对应的后端即可。提示如果你也想自己写一个框架我强烈建议第一版不要碰 GPU。先把 CPU 版本的自动求导跑通理解清楚计算图的构建和反向传播的链路再去考虑硬件加速的事情。3. 核心模块的详细设计与实操要点3.1 张量模块不只是 NumPy 的包装张量是框架中最基础的数据结构。它需要存储两部分信息数据本身和元信息。数据就是一个多维数组元信息包括形状、数据类型、是否需要梯度等。我设计的张量类大致长这样内部用一个 NumPy 数组来存实际数据同时维护一个requires_grad标志来表示这个张量是否需要计算梯度。如果requires_grad为 True还需要一个grad属性来存储反向传播时累积的梯度值。这里有一个容易忽略的细节张量的运算结果是否需要梯度取决于它的输入中是否有需要梯度的张量。比如c a b如果 a 需要梯度而 b 不需要那么 c 也需要梯度。这个判断逻辑必须写在每一个运算函数里否则反向传播链就会断掉。另一个关键点是计算历史的记录。每次对张量做运算时除了返回结果还需要记录“这个结果是由哪个操作、哪些输入产生的”。我采用的方式是在张量对象上挂一个_ctx属性里面保存操作类型和输入张量的引用。这样在反向传播时就可以沿着这个引用链一路回溯。class Tensor: def __init__(self, data, requires_gradFalse): self.data np.array(data, dtypenp.float32) self.requires_grad requires_grad self.grad None self._ctx None # 记录计算历史 def __add__(self, other): other other if isinstance(other, Tensor) else Tensor(other) out Tensor(self.data other.data, requires_gradself.requires_grad or other.requires_grad) if out.requires_grad: out._ctx (add, self, other) return out上面这段代码是整个框架最核心的几行之一。它展示了张量运算的基本模式计算结果、判断是否需要梯度、记录计算历史。后面的所有运算包括乘法、矩阵乘法、激活函数都遵循同样的模式。注意在记录计算历史时一定要保存输入张量的引用而不是只保存数据。因为反向传播时需要用到输入张量的值来计算梯度如果只保存数据副本不仅浪费内存还可能因为数据被修改而导致梯度计算错误。3.2 自动求导模块反向传播的引擎自动求导的核心思想是链式法则。如果y f(x)那么dy/dx f(x)。如果z g(y)那么dz/dx g(y) * f(x)。在计算图中每个节点只需要知道如何计算自己相对于输入的局部梯度然后把这些局部梯度沿着图反向传播回去就行了。我在实现时采用了一种“反向拓扑排序”的策略。具体来说从最终输出节点开始按照计算历史的逆序依次访问每个节点。对于每个节点根据它记录的操作类型调用对应的反向函数计算出输入张量的梯度并累积到输入张量的grad属性上。这里有一个关键问题梯度累积。如果一个张量被多个后续操作使用那么它的梯度应该是所有路径传回来的梯度之和。比如y x * 2和z x * 3那么x的梯度应该是dy/dx dz/dx 2 3 5。所以在反向传播时每次计算出梯度后不是直接赋值而是累加。def backward(self, gradNone): if grad is None: grad np.ones_like(self.data) if self.grad is None: self.grad grad else: self.grad grad if self._ctx is None: return op, left, right self._ctx if op add: if left.requires_grad: left.backward(grad) if right.requires_grad: right.backward(grad) elif op mul: if left.requires_grad: left.backward(grad * right.data) if right.requires_grad: right.backward(grad * left.data) # 其他操作类似上面这段代码展示了反向传播的基本框架。对于加法操作梯度直接传递给两个输入对于乘法操作梯度需要乘以另一个输入的值。这就是链式法则在计算图上的具体体现。实操心得在实现反向传播时最容易出错的地方是梯度的形状。比如矩阵乘法C A BA 的形状是 (m, n)B 的形状是 (n, p)C 的形状是 (m, p)。那么dA dC B.TdB A.T dC。如果形状对不上NumPy 会报错或者产生错误的结果。我建议每实现一个操作的反向传播都先用数值梯度验证一下。3.3 计算图模块动态构建与内存管理在动态图模式下计算图不是预先定义的而是随着前向传播的进行逐步构建的。每次运算都会创建一个新的节点节点之间通过输入输出关系连接。这种“边运行边建图”的方式使得图的结构完全由实际的代码执行路径决定。这样做的一个直接后果是每次前向传播都会构建一张新的计算图。对于训练循环来说这意味着每处理一个 batch 的数据都会生成一张新的图。这听起来很浪费但实际上正是动态图的灵活性所在。因为不同的 batch 可能有不同的序列长度、不同的分支路径静态图很难处理这种变化。内存管理是计算图模块中另一个需要认真考虑的问题。在前向传播过程中每个中间结果都需要保存下来因为反向传播时要用到。如果模型很深中间结果会占用大量内存。PyTorch 的做法是在反向传播过程中逐步释放不再需要的中间结果我在框架中也采用了类似的策略当一个节点的梯度计算完成后如果它的所有下游节点都已经处理完毕就可以释放它的数据。不过在第一版中我并没有实现复杂的内存回收机制而是简单地依赖 Python 的垃圾回收。因为对于学习性质的框架来说先把逻辑跑通比优化内存更重要。如果你后面想把这个框架用于实际训练内存管理是一个必须认真对待的问题。3.4 神经网络层模块从张量到模型有了张量和自动求导搭建神经网络层就变成了一个相对直接的事情。一个全连接层本质上就是y x W b其中 W 和 b 是需要梯度的参数。激活函数如 ReLU 就是y max(0, x)它的反向传播也很简单如果输入大于 0梯度直接传递否则梯度为 0。我在设计层模块时采用了一个基类Layer它定义了forward和parameters两个接口。所有具体的层都继承这个基类实现自己的前向逻辑并返回需要优化的参数列表。这样做的好处是后面写优化器时只需要遍历模型的parameters()就能拿到所有需要更新的张量不需要关心模型的具体结构。class Linear(Layer): def __init__(self, in_features, out_features): self.W Tensor(np.random.randn(in_features, out_features) * 0.01, requires_gradTrue) self.b Tensor(np.zeros(out_features), requires_gradTrue) def forward(self, x): return x self.W self.b def parameters(self): return [self.W, self.b]上面这个全连接层的实现非常简洁但它已经包含了神经网络层的核心要素参数初始化、前向计算、参数暴露。损失函数如均方误差MSE mean((y_pred - y_true)^2)也可以用同样的方式实现只需要定义好前向计算和对应的反向传播逻辑。提示参数初始化对训练效果影响很大。全连接层的权重我通常用小的随机数初始化比如乘以 0.01避免初始输出过大导致梯度爆炸。偏置项初始化为 0 是常见做法。这些细节在写框架时看起来不起眼但实际训练时差别很明显。3.5 优化器模块参数更新的策略优化器的职责很简单拿到梯度按照某种规则更新参数。最基础的随机梯度下降SGD就是param param - lr * grad。虽然简单但它是理解更复杂优化器的基础。我在实现优化器时把它设计成一个独立的类接收模型的参数列表和学习率作为输入。每次调用step()方法时遍历所有参数根据它们的grad属性进行更新。然后调用zero_grad()清空梯度为下一次迭代做准备。class SGD: def __init__(self, params, lr0.01): self.params params self.lr lr def step(self): for p in self.params: if p.grad is not None: p.data - self.lr * p.grad def zero_grad(self): for p in self.params: p.grad None这段代码虽然短但它揭示了优化器的本质优化器不关心模型结构只关心参数和梯度。这种解耦设计使得你可以轻松地替换不同的优化器比如后面要实现的 Momentum、Adam 等只需要修改参数更新的公式即可。4. 实操过程与核心环节实现4.1 从零搭建一个可训练的全连接网络理论说了这么多最终还是要落到代码上。我以一个简单的两层全连接网络为例展示整个框架是如何协同工作的。这个网络的任务是拟合一个非线性函数比如y sin(x)。首先定义模型结构输入维度 1隐藏层维度 16输出维度 1。隐藏层使用 ReLU 激活输出层不加激活。损失函数用均方误差。class TwoLayerNet: def __init__(self): self.fc1 Linear(1, 16) self.fc2 Linear(16, 1) def forward(self, x): h self.fc1.forward(x) h relu(h) return self.fc2.forward(h) def parameters(self): return self.fc1.parameters() self.fc2.parameters()然后写训练循环。每个迭代做四件事前向传播计算预测值计算损失反向传播计算梯度优化器更新参数。model TwoLayerNet() optimizer SGD(model.parameters(), lr0.01) for epoch in range(1000): x Tensor(np.random.uniform(-3, 3, (32, 1))) y_true Tensor(np.sin(x.data)) y_pred model.forward(x) loss mse_loss(y_pred, y_true) loss.backward() optimizer.step() optimizer.zero_grad() if epoch % 100 0: print(fEpoch {epoch}, Loss: {loss.data:.4f})这个训练循环看起来和 PyTorch 的写法非常相似这不是巧合。因为 PyTorch 的核心设计理念就是“让用户像写普通 Python 代码一样写模型”而我在设计这个框架时也遵循了同样的原则。通过亲手实现这些逻辑我对loss.backward()背后发生了什么有了完全不同的理解。4.2 梯度验证确保反向传播的正确性写完反向传播之后最重要的一步是验证它的正确性。我用的方法是数值梯度检验。对于任意一个参数它的数值梯度可以通过有限差分来近似(f(x eps) - f(x - eps)) / (2 * eps)。把这个数值梯度和反向传播计算出的解析梯度对比如果误差在 1e-5 以内就说明实现是正确的。def numerical_gradient(f, x, eps1e-5): grad np.zeros_like(x.data) it np.nditer(x.data, flags[multi_index]) while not it.finished: idx it.multi_index old_val x.data[idx] x.data[idx] old_val eps f_pos f().data.sum() x.data[idx] old_val - eps f_neg f().data.sum() grad[idx] (f_pos - f_neg) / (2 * eps) x.data[idx] old_val it.iternext() return grad这个检验方法虽然计算量大但对于小规模的网络来说完全可行。我在开发过程中每实现一个新的操作都会用这个方法验证一遍。有几次就是靠它发现了反向传播中的符号错误和形状错误。实操心得梯度检验时建议把网络调小比如输入输出都设为 2 维隐藏层设为 3 维。这样数值梯度的计算量可以接受而且一旦出错定位问题也更容易。另外记得在检验前把参数随机种子固定保证每次运行的结果可复现。4.3 计算图的构建过程实录为了更直观地理解动态图的构建过程我以y (x1 x2) * x3这个表达式为例追踪一下计算图是如何一步步建立的。第一步x1 x2触发Tensor.__add__创建一个新的张量t1其_ctx记录为(add, x1, x2)。第二步t1 * x3触发Tensor.__mul__创建张量y其_ctx记录为(mul, t1, x3)。此时计算图的结构是x1 和 x2 指向 t1t1 和 x3 指向 y。反向传播时从 y 开始。y 的梯度初始化为 1。根据_ctx中的mul操作t1 的梯度为1 * x3.datax3 的梯度为1 * t1.data。然后 t1 继续反向传播根据add操作x1 和 x2 的梯度都等于 t1 的梯度。整个过程完全由_ctx中记录的信息驱动不需要任何额外的图结构。这种设计的美妙之处在于计算图不是显式存储的而是隐含在张量的引用关系中。每个张量只需要知道自己是怎么来的反向传播就能自动完成。这也是动态图框架代码量少、逻辑清晰的根本原因。4.4 性能优化的几个实用技巧虽然这个框架的主要目的是学习但在实现过程中我还是做了一些性能优化让它在小规模任务上不至于太慢。第一个优化是避免不必要的张量创建。在反向传播时很多中间梯度其实不需要保存为完整的张量对象直接用 NumPy 数组计算即可。我在反向传播的内部实现中尽量使用 NumPy 的原地操作减少内存分配。第二个优化是批处理矩阵运算。全连接层的前向传播x W b中x 是一个 batch 的数据矩阵乘法会自动利用 NumPy 的底层优化。这比逐个样本计算要快得多。第三个优化是延迟梯度清零。标准的做法是每次迭代后调用zero_grad()但实际上可以在下一次反向传播时再清零这样可以避免一次不必要的遍历。不过为了代码清晰我还是保留了显式的zero_grad()。优化手段效果实现难度原地操作减少内存分配低批处理矩阵运算利用 BLAS 加速低延迟梯度清零减少一次遍历中中间结果释放降低内存峰值高上面这张表总结了我在第一版中考虑过的优化手段。对于学习性质的框架前两项就足够了。后两项如果要做需要更精细的内存管理容易引入 bug建议在框架稳定之后再考虑。5. 常见问题与排查技巧实录5.1 梯度为 None 或形状不匹配这是最常见的问题。表现是训练完全不收敛或者报错说grad是 None。根本原因通常是计算链在某处断了。比如某个操作没有正确设置requires_grad或者反向传播函数中没有把梯度传递给某个输入。排查方法是逐层打印梯度的形状和值。从损失函数开始往前检查每一层的梯度是否存在、形状是否正确。如果某一层的梯度是 None就检查它的前向传播是否记录了计算历史以及反向传播是否被正确调用。另一个常见原因是原地操作破坏了计算图。比如x.data 1这种操作会直接修改张量的数据但不会记录到计算历史中。如果这个张量后续参与了需要梯度的运算反向传播时就会用到被修改后的数据导致梯度错误。正确的做法是使用x x 1这样会创建一个新的张量并记录计算历史。5.2 损失值变成 NaN训练过程中损失突然变成 NaN通常意味着梯度爆炸或者出现了除零操作。可能的原因包括学习率太大、参数初始化不当、损失函数中有 log(0) 或除以零的情况。我的排查步骤是这样的首先把学习率调小一个数量级看是否还会出现 NaN。如果问题依旧检查损失函数的实现确保没有对零取对数或者除以零。然后在反向传播的每个节点打印梯度的最大值和最小值找到梯度异常的那一层。注意ReLU 激活函数在输入为负时梯度为 0这本身没问题。但如果某一层的输出全部为负那么这一层的梯度就全部为 0参数永远不会更新。这种情况叫做“神经元死亡”。解决办法是使用 LeakyReLU 或者调整初始化方式让初始输出有正有负。5.3 训练速度异常缓慢如果发现训练速度比预期慢很多可以从几个方面排查。第一检查是否在每次迭代中都重新创建了计算图而没有释放旧图导致内存持续增长。第二检查是否有 Python 层面的循环拖慢了矩阵运算比如逐元素操作而不是批量操作。第三检查 NumPy 是否链接了优化的 BLAS 库比如 OpenBLAS 或 MKL这会让矩阵乘法的速度差好几倍。还有一个容易被忽略的点是数据类型。如果张量默认使用 float64计算量会是 float32 的两倍。我在框架中统一使用 float32既节省内存又加快计算。你可以在创建张量时显式指定dtypenp.float32。问题现象可能原因排查方法梯度为 None计算链断裂逐层打印梯度损失 NaN梯度爆炸/除零调小学习率检查损失函数训练缓慢内存泄漏/数据类型检查内存增长统一 float32不收敛初始化不当/学习率梯度检验调整超参数这张表可以作为日常调试的速查参考。实际遇到的问题可能更复杂但大多数情况下都能从这几个方向找到线索。5.4 反向传播中的循环引用问题在实现计算图时如果处理不当可能会产生循环引用。比如两个张量互相引用对方的_ctx导致 Python 的垃圾回收器无法释放它们。这个问题在长时间训练中会逐渐暴露出来表现为内存持续增长。我的解决办法是在反向传播完成后手动断开计算历史中的引用。具体来说在backward()函数的最后把当前节点的_ctx设为 None。这样当反向传播走到图的起点时整条链上的引用都会被释放。虽然这样做会让计算图无法被重复使用但对于训练场景来说每次迭代都会构建新的图旧图本来就不需要保留。提示如果你想让计算图支持多次反向传播比如计算二阶梯度那就不能断开引用。这种情况下需要更复杂的内存管理策略比如使用弱引用或者显式的图生命周期管理。第一版框架不建议做这么复杂。5.5 参数更新后梯度没有清零这是一个很隐蔽的 bug。如果你忘记调用zero_grad()梯度会在每次反向传播时累加导致参数更新量越来越大训练很快就会发散。表现是损失先下降然后突然爆炸。排查方法很简单在每次optimizer.step()之后打印几个参数的梯度值看看它们是否在每次迭代后归零。如果没有归零就检查zero_grad()是否被正确调用。我在框架中把zero_grad()设计成优化器的方法而不是模型的方法就是为了让用户不容易忘记——因为优化器是必须显式创建和调用的把清零逻辑放在这里更符合直觉。6. 后续扩展方向与个人体会这个框架的第一版到这里就基本成型了。它虽然简单但已经具备了深度学习框架的核心要素张量运算、自动求导、计算图、神经网络层和优化器。用这个框架训练一个小型的全连接网络拟合简单的非线性函数完全没有问题。后续如果要继续扩展我觉得有几个方向值得尝试。第一个是卷积层和池化层的实现这是处理图像数据的基础。卷积的反向传播比全连接层复杂一些但原理是一样的都是链式法则的应用。第二个是更复杂的优化器比如 Momentum 和 Adam它们的实现只需要在 SGD 的基础上增加一些状态变量。第三个是计算图的优化比如算子融合和内存复用这需要对计算图做静态分析是一个很有挑战性的方向。我个人在实际操作中的体会是写框架最大的收获不是框架本身而是对深度学习底层原理的理解。以前用 PyTorch 的时候loss.backward()就是一个黑盒现在我知道它背后发生了什么知道梯度是如何一步步传回去的知道为什么有些操作会导致梯度消失或爆炸。这种理解在调试复杂模型时非常有用因为你能从原理出发去分析问题而不是盲目地试错。最后再分享一个小技巧如果你也想写自己的框架建议从最简单的标量自动求导开始比如实现一个支持加减乘除的标量类然后逐步扩展到张量。这样每一步的复杂度都可控遇到问题也容易定位。我一开始就想直接写张量版本结果在形状处理上卡了很久后来退回去从标量做起反而更快地理解了核心逻辑。