从零手写深度学习框架：Tensor、计算图与Autograd架构设计

1. 为什么我要从零写一个深度学习框架动手写一个深度学习框架这件事我在脑子里盘算了很久。起因很简单用现成的框架做模型训练调包调了几年model.fit()一按就跑但真遇到梯度爆炸、loss 不收敛、显存莫名爆掉的时候心里其实是虚的。你知道大概该调学习率、该加正则但底层到底发生了什么autograd 是怎么把梯度一层层传回去的优化器里的动量到底怎么更新的这些细节长期被封装在黑盒里。所以这个系列我打算干一件事从零开始用纯 Python 加 NumPy搭一个能跑通前向传播、反向传播、参数更新的迷你深度学习框架。这一篇是第 1 篇只讲架构设计不写具体实现代码把整个框架的骨架、模块划分、数据流向先定下来。后面几篇再逐个模块填肉。这篇文章适合谁看如果你已经会用 PyTorch 或 TensorFlow 训练模型但想搞清楚框架内部到底怎么运转的那正好。如果你刚学完反向传播的数学推导想找个工程落地的抓手也合适。纯小白可能需要先补一下链式法则和矩阵求导但我会尽量把每个设计决策背后的“为什么”讲清楚让你知其然也知其所以然。我先把结论摆出来一个能用的深度学习框架核心就四块——张量Tensor作为数据载体、计算图Computation Graph记录运算关系、自动微分Autograd反向求梯度、优化器Optimizer更新参数。把这四块的关系理顺框架就立起来了。剩下的损失函数、层封装、数据加载都是在这四块之上的应用层。2. 框架整体架构与模块划分2.1 四层架构的划分逻辑我在设计的时候把整个框架分成四层从下往上依次是基础层、计算层、网络层、应用层。这个分层不是拍脑袋定的而是照着“数据怎么流、梯度怎么回、参数怎么更新”这条主线切出来的。基础层放的是最底层的数学工具比如矩阵乘法、广播机制、随机数生成。这一层不涉及任何“学习”的概念纯粹是数值计算。计算层是框架的心脏包含 Tensor 对象、计算图节点、自动微分引擎。网络层负责把计算层的算子组装成层Layer和模型Model提供forward接口。应用层就是损失函数、优化器、训练循环这些面向使用者的东西。为什么要这么分因为分层能让每一层的职责单一测试和替换都方便。比如我想把底层的 NumPy 换成自己写的矩阵运算库只要基础层接口不变上面三层完全不用动。再比如我想加一个新的优化器只在应用层加就行不碰下面的计算逻辑。这种解耦在框架开发里特别重要否则代码会迅速腐化成意大利面条。提示分层设计的一个实用判断标准是——如果修改某一层的代码导致另一层必须跟着改说明分层边界没划对。我在早期版本里把优化器的更新逻辑写进了 Tensor 类里结果每加一个优化器都要改 Tensor后来拆出来才清爽。2.2 核心模块清单与依赖关系具体到模块我列了一张清单每个模块的职责和依赖关系如下表。这张表是整个框架的“施工图”后面写代码基本就是照着它逐个实现。模块名称所属层级核心职责依赖模块Tensor计算层数据存储、形状管理、运算重载基础层Function计算层定义前向/反向计算构建计算图TensorAutograd计算层拓扑排序、梯度回传调度Function、TensorLayer网络层参数管理、前向封装Tensor、FunctionModel网络层层的容器、参数收集LayerLoss应用层损失计算及梯度起点Tensor、FunctionOptimizer应用层参数更新策略TensorDataLoader应用层批次划分、数据迭代基础层依赖关系是单向的上层依赖下层下层不知道上层的存在。这一点很关键。Tensor 不应该知道优化器是什么Function 也不应该知道模型是什么。依赖单向保证了框架的可扩展性你加新功能时只需要往上叠不用回头改底层。2.3 数据流向与梯度回传路径整个框架的数据流可以概括成一句话前向建图反向回传最后更新。前向传播时每个运算都会创建一个 Function 节点记录输入 Tensor 和输出 Tensor同时把节点挂到计算图上。比如y x * w b这里会生成两个节点一个乘法节点、一个加法节点它们通过 Tensor 的creator指针串成一条链。反向传播时从损失值出发沿着计算图反向遍历每个节点调用自己的backward方法把上游传来的梯度乘以本地的局部梯度再传给下游。这里有个设计选择值得说计算图是动态构建的还是静态构建的。静态图像早期的 TensorFlow先定义再运行动态图像 PyTorch边运行边建图。我选动态图原因是调试友好——你可以在任意位置打断点看中间结果出错了堆栈也清晰。静态图虽然理论上能做更多图优化但对一个教学性质的框架来说可调试性比性能优化重要得多。梯度回传的路径和前向路径正好相反但节点顺序需要保证一个节点必须等它所有的下游节点都把梯度传回来之后才能计算自己的梯度。这就是为什么反向传播前要做一次拓扑排序。如果图里有分支一个 Tensor 被多个运算使用它的梯度是各分支梯度之和这个累加逻辑必须在 Autograd 里处理干净否则梯度会算错。3. 核心模块的深度拆解与设计取舍3.1 Tensor 设计数据载体还是计算节点Tensor 是整个框架里最纠结的一个类。它到底应该只是“一块带形状的数据”还是同时承担“计算图节点”的角色我最后的选择是让 Tensor 同时承担这两个职责但把计算逻辑委托给 Function。具体来说Tensor 内部存三样东西dataNumPy 数组存数值、grad梯度形状和 data 一致、creator指向创建它的 Function 节点。当两个 Tensor 做运算时__add__、__mul__这些魔术方法被触发它们不直接算而是调用对应的 Function 实例由 Function 完成计算并返回新的 Tensor同时把新 Tensor 的creator指向自己。为什么不让 Tensor 自己算因为把计算逻辑和数据结构分离能让反向传播的实现更统一。每个 Function 只需要实现forward和backward两个方法Autograd 引擎遍历图时统一调用backward不用关心具体是什么运算。如果计算逻辑散落在 Tensor 的各个方法里反向传播就得写一堆 if-else 判断运算类型维护起来是灾难。注意Tensor 的grad默认是 None只有在反向传播经过它时才会被分配。这个懒分配策略能省不少内存尤其是中间激活值很多根本不需要梯度。3.2 计算图构建动态图的节点管理计算图的节点就是 Function 实例。每个 Function 在 forward 时接收若干输入 Tensor产出输出 Tensor并把输入存下来供 backward 使用。节点之间的连接靠 Tensor 的creator指针和 Function 的inputs列表双向维护。我举个具体例子说明节点怎么串起来。假设做z relu(x w b)运算顺序是先x w得到t1再t1 b得到t2最后relu(t2)得到z。对应的节点链是MatMul节点产出t1Add节点产出t2ReLU节点产出z。z.creator是 ReLU 节点ReLU 节点的输入是t2t2.creator是 Add 节点以此类推一直回溯到叶子节点也就是参数和输入数据。反向传播时从z出发先调 ReLU 的 backward算出t2的梯度再调 Add 的 backward算出t1和b的梯度最后调 MatMul 的 backward算出x和w的梯度。整个过程就是沿着 creator 链一路回溯。这里有个坑我踩过如果同一个 Tensor 被用了两次梯度必须累加而不是覆盖。比如y x xx 的梯度应该是 2 而不是 1。解决办法是在 backward 里用而不是并且在每次反向传播开始前把图里所有 Tensor 的 grad 清零。这个清零操作放在优化器的zero_grad里做和 PyTorch 的设计一致。3.3 Autograd 引擎拓扑排序与梯度调度Autograd 引擎的核心任务就两个确定反向遍历的顺序以及调度每个节点的 backward。顺序问题用拓扑排序解决调度问题用递归或迭代实现。拓扑排序的思路是从损失 Tensor 出发做一次深度优先搜索记录每个节点的完成时间最后按完成时间逆序排列就得到反向传播的顺序。为什么是这个顺序因为拓扑排序保证了一个节点在它所有依赖节点之前被处理而反向传播要求下游先于上游所以逆序正好。我用一个简单的例子说明。图是a - b - c正向是 a 算出 bb 算出 c。拓扑排序从 c 出发DFS 访问 b再访问 aa 完成、b 完成、c 完成逆序就是 c、b、a正好是反向传播需要的顺序。调度实现上我选迭代而非递归。递归写起来简洁但图深了会爆栈尤其是深层网络。迭代版本用一个列表当栈手动管理访问状态虽然代码长一点但稳。这个取舍在框架开发里很常见教学代码可以递归生产代码必须迭代。提示拓扑排序时要注意处理“一个节点被多次访问”的情况。用一个 visited 集合记录已处理的节点避免重复入栈否则图有共享子结构时会做无用功。3.4 参数管理与优化器接口设计参数在框架里就是“需要梯度的 Tensor”。我在 Tensor 上加了一个requires_grad标志只有为 True 的 Tensor 才会在反向传播时被分配梯度。Layer 在初始化时创建参数 Tensor并统一注册到一个参数列表里Model 再把所有 Layer 的参数收集起来交给优化器。优化器的接口设计我参考了主流框架的约定核心就三个方法zero_grad清空梯度、step更新参数、state_dict保存状态。step里做的是param.data - lr * param.grad如果是带动量的优化器还要维护一个速度变量。为什么把优化器和参数分开因为同一个优化器可以管多组参数同一组参数也可以换不同优化器。这种灵活性在调参时很有用比如你可以给不同的层设不同的学习率。如果把更新逻辑写死在参数里这种灵活性就没了。优化器更新公式适用场景SGDp p - lr * g简单任务、需要精细调参Momentumv m * v - lr * g; p v震荡严重的损失面Adam一阶二阶动量自适应大多数默认场景4. 实操从零搭建框架骨架的完整流程4.1 环境准备与目录结构动手之前先把环境理清楚。我用的就是最朴素的组合Python 3.10 加 NumPy不装任何深度学习相关的库这样才能保证框架是“从零”的。NumPy 只用来做底层矩阵运算autograd 和计算图全部自己写。目录结构我按模块划分每个模块一个文件方便后续逐个实现和测试mini-dl/ ├── core/ │ ├── tensor.py # Tensor 类 │ ├── function.py # Function 基类及各类算子 │ └── autograd.py # 自动微分引擎 ├── nn/ │ ├── layer.py # Layer 基类 │ ├── linear.py # 全连接层 │ └── model.py # Model 容器 ├── optim/ │ ├── sgd.py # SGD 优化器 │ └── adam.py # Adam 优化器 ├── loss/ │ └── mse.py # 均方误差 └── tests/ └── test_autograd.py这个结构的好处是每个文件职责单一测试时能单独 import。比如我想验证 autograd 对不对直接跑test_autograd.py不用把整个框架跑起来。4.2 关键接口定义与参数计算在写实现之前我先把几个核心接口的签名定下来这样后面填代码时不会来回改。Tensor 的构造函数接收 data 和 requires_grad内部把 data 转成 NumPy 数组并记录形状。Function 基类定义forward和backward两个抽象方法__call__里负责调用 forward、包装输出 Tensor、设置 creator。Autograd 引擎暴露一个backward函数接收损失 Tensor内部做拓扑排序和梯度调度。参数初始化这块有个细节值得算一下。全连接层的权重我用 Xavier 初始化公式是std sqrt(2 / (fan_in fan_out))。为什么用这个因为如果权重初始化太大前向传播时激活值会逐层放大最后溢出太小则逐层缩小梯度消失。Xavier 让每一层的输出方差和输入方差保持一致信号能稳定传下去。我实测过用 Xavier 初始化的网络比用纯随机初始化的收敛快不少尤其是层数多的时候。import numpy as np def xavier_init(fan_in, fan_out): std np.sqrt(2.0 / (fan_in fan_out)) return np.random.randn(fan_in, fan_out) * std4.3 前向传播与反向传播的联调骨架搭好后第一件事是验证前向和反向能不能对上。我的做法是用数值梯度校验解析梯度对一个简单函数分别用 autograd 算梯度和用有限差分算梯度两者应该非常接近。数值梯度的公式是(f(xh) - f(x-h)) / (2h)h 取 1e-5 左右。这个校验方法在框架开发里是标配每次加新算子都要跑一遍。我踩过的坑是如果 h 取太大截断误差大取太小浮点误差大。1e-5 是个经验值实测下来对 float64 精度够用。联调时我建议从最简单的运算开始先验证加法、乘法再验证矩阵乘法最后验证完整的线性层。每加一个算子就测一个不要攒一堆再测否则出了问题很难定位是哪个算子的锅。注意数值梯度校验只在 float64 下可靠float32 精度不够误差会大到无法判断。所以测试时统一用 float64训练时再切回 float32。4.4 一个最小可跑通的训练闭环骨架验证通过后我搭了一个最小闭环一个线性层加 MSE 损失用 SGD 优化拟合一个简单的线性关系。这个闭环虽然简单但把前向、反向、更新三个环节全串起来了能跑通就说明框架的基本逻辑没问题。训练循环长这样取一批数据前向算预测值算损失反向传播优化器更新清空梯度。这五步是所有深度学习训练的通用骨架后面加再多花哨的东西本质还是这五步。我建议你把这个循环背下来调任何模型都是这个套路。跑通之后我观察了一下 loss 曲线前几十步下降很快后面趋于平缓这是正常的收敛形态。如果 loss 不降反升八成是学习率太大或者梯度符号搞反了如果 loss 几乎不动可能是初始化有问题或者梯度没传回去。这些排查经验后面会专门讲。5. 常见问题与排查技巧实录5.1 梯度相关问题的排查思路梯度问题是框架开发里最高频的坑我整理了一张速查表遇到问题按表排查能省不少时间。现象可能原因排查方法梯度全为 0计算图断了某节点没设 creator打印每个 Tensor 的 creator 是否为 None梯度爆炸学习率太大或初始化太大打印梯度范数看是否逐层放大梯度消失激活函数饱和或初始化太小检查激活值分布换 ReLU梯度形状不对backward 里维度没对齐对比 grad 和 data 的 shape梯度累加错误共享 Tensor 的梯度没累加检查是否用了 而非 排查梯度问题我有个笨办法但很管用在 backward 里加打印把每个节点的输入梯度、输出梯度、局部梯度都打出来。虽然输出多但一眼就能看出哪一层开始不对劲。等定位到问题再删掉打印。还有一个隐蔽的坑broadcasting 的反向传播。前向时如果发生了广播比如 (3,4) 和 (4,) 相加反向时梯度需要沿着广播的维度求和还原。这个逻辑如果写错梯度形状就对不上报错信息还特别难懂。我的建议是广播的反向单独写测试别和普通运算混在一起测。5.2 计算图构建的典型陷阱计算图这块我踩过两个印象深刻的坑。第一个是原地操作。如果我在 Tensor 上做x 1这种原地修改会破坏计算图记录的前向值导致反向传播时用的值和前向时不一致梯度算错。解决办法是所有运算都返回新 Tensor绝不原地修改。这个约束在框架里要严格执行宁可多占点内存。第二个是图的重复构建。如果我在训练循环里忘了清空计算图每次迭代的节点会不断累积内存越吃越多最后 OOM。动态图框架一般靠 Python 的垃圾回收自动清理但前提是没有循环引用。我在 Function 里存输入 Tensor 时用的是弱引用避免节点之间互相持有导致回收不掉。提示调试计算图时可以写一个visualize_graph函数从损失 Tensor 出发遍历 creator 链把节点类型和形状打印成树状结构。这个工具在排查“图断在哪”时特别好用。5.3 性能与内存的优化经验框架能跑通之后下一步就是让它跑得快、吃得少。我总结了几个实用的优化点。第一减少不必要的 Tensor 创建。每次运算都新建 Tensor 是有开销的尤其是小运算频繁调用时。我后来加了一个内存池复用形状相同的 Tensor 的底层数组实测能省 20% 左右的内存分配时间。第二反向传播时及时释放中间梯度。不是所有中间 Tensor 的梯度都需要保留只有参数的梯度需要。我在 Autograd 里加了个标记非参数的中间梯度算完就置 None让 GC 回收。第三矩阵运算尽量批量做。框架层面的循环能避免就避免把 batch 维度利用起来让 NumPy 走底层优化过的 BLAS 库。我试过把逐样本的循环改成批量矩阵乘速度差了几十倍。这些优化不是一开始就要做的先保证正确再谈性能。我见过太多人一上来就抠性能结果正确性都没保证优化了半天全是错的。正确的顺序是跑通、测对、再优化。5.4 从教学框架到可用框架的差距最后说点实在的。这个系列搭出来的框架定位是教学和理解原理不是替代 PyTorch。真要用到生产里还差不少东西GPU 支持、算子融合、自动混合精度、分布式训练、完整的算子库、成熟的调试工具。这些每一个都是大工程。但理解了这个迷你框架你再看 PyTorch 的源码就不会发怵了。它的核心机制——Tensor、autograd、计算图、优化器——和这里讲的是一回事只是工程实现上做了大量优化和扩展。先在小框架里把原理吃透再去读大框架的源码这条路我觉得是最扎实的。我自己在实际操作中的体会是写框架最大的收获不是框架本身而是对“梯度是怎么流动的”这件事有了肌肉记忆。以前调模型靠试现在能根据 loss 曲线和梯度范数判断问题出在哪一层。这种直觉光看文档是练不出来的必须自己动手写一遍。后面几篇我会逐个模块把代码填上从 Tensor 开始一步步把这个骨架变成能跑的东西。