零依赖手写深度学习框架:自动求导与计算图实践 第一次把一套手写的训练脚本跑出正常下降的损失曲线是在修好广播梯度归约之后的第三天。前三天里损失函数像被钉在 2.30 附近不动我一度怀疑是学习率、初始化、甚至是随机种子的锅最后发现是(N, C)的张量加(C,)的偏置时反向传播只把梯度传回了一半的形状剩下一半被默默丢掉了。这个坑后来成了我拉人入伙时必讲的反面教材也是我给这套东西起名DeepZero的原因——把深度学习栈的外部依赖降到零从张量、计算图、反向传播到优化器一环一环自己写一遍写到你亲眼看见梯度是怎么流的。DeepZero 不是什么生产级框架它是一个教学与验证性质的自研项目目标很明确用一个几百行的核心库把前向构图—反向求导—参数更新这条链路完整跑通能训练一个两层全连接网络在二维分类任务上收敛也能手写一个小型序列模型跑通字符级预测。它适合三类人刚学完反向传播公式但对代码实现没底的在校同学、想搞懂框架黑盒内部到底在算什么的后端或数据方向开发者、以及需要给模型做定制算子但被自动求导挡住的工程同学。下面的内容全部围绕这套实现展开包含字段设计、形状对齐、梯度校验、优化器状态管理和几个我真实踩过的坑代码可以直接复现。1. DeepZero 的定位零依赖不等于零工程先说清楚零依赖的边界否则这个项目会失控。我的定义是深度学习的核心机制零依赖工程配套允许外包。也就是说张量存储、计算图、自动求导、优化器更新、损失函数这五块必须手写一行现成的自动求导都不能调而底层的数组运算矩阵乘法、广播、转置交给 NumPy画损失曲线交给 Matplotlib数据加载用最朴素的列表切片。这个切分不是偷懒而是有明确理由的底层 BLAS 层面的矩阵乘法优化涉及分块、缓存命中、SIMD 指令属于高性能计算领域跟理解深度学习怎么训练这件事没有因果关系自己写只会得到比 NumPy 慢两个数量级的结果还会掩盖真正的学习点。真正的学习点在计算图怎么记录、梯度怎么回传、形状怎么对齐这三件事上而这三件事恰恰是任何成熟框架都不会在文档里讲透的部分。1.1 为什么不建议照抄成熟框架的目录结构我最初的做法是照着印象里的框架源码结构建目录tensor/、autograd/、nn/、optim/、functional/、utils/一共七个包二十多个文件。结果是两周之后我连自己写的functional和nn里哪份是重复实现都分不清了。后来我把整个项目砍到三个文件core.pyTensor 与计算图、nn.py层与损失、train.py训练循环代码量从两千多行降到六百行左右反而在一个下午就跑通了。原因很直白目录结构的复杂度会伪装成项目的复杂度。当你有七个包的时候你会觉得我已经完成了架构设计实际上核心的backward函数一行都没写。而当你只有三个文件的时候你没法逃避那个必须面对的 30 行反向传播实现。所以 DeepZero 的目录原则是在你写完全部核心机制、并且用手算数值验证过梯度正确之前不允许拆分模块。拆分是重构动作不是设计动作。1.2 哪些东西必须自己啃哪些坚决外包这张表是我反复调整后的最终划分供参考模块处理方式理由张量存储与形状自己封装 NumPy 数组需要挂grad、_backward等元信息计算图记录与前向传播必须手写理解动态图的核心反向传播链式法则必须手写全项目唯一不可替代的学习点广播梯度归约必须手写最高频的静默错误来源优化器状态更新必须手写理解动量、二阶矩的维护逻辑矩阵乘法内核调用 NumPy属于高性能计算与训练机制无关数据加载与打乱朴素实现即可不属于核心链路可视化调用 Matplotlib纯展示数值梯度校验必须手写验证正确性的唯一可靠手段这里有个容易忽略的点数值梯度校验必须自己写而且要写得比反向传播更仔细。因为它是你唯一的真值来源如果校验代码本身有 bug你会得到一个虚假的通过或者虚假的失败然后花大量时间去修一个根本不存在的问题。我就经历过一次校验器写错方向的情况把(f(xh) - f(x-h)) / 2h写成了(f(xh) - f(x)) / h导致所有梯度都显示有千分之一量级的误差我花了整整一晚去排查一个精度完全正常的实现。2. Tensor 的字段设计够用就好多一个都是负债Tensor 是整个项目的地基它的字段设计直接决定了后面所有代码的写法。DeepZero 里的 Tensor 只有五个字段我把它们列出来并逐个解释为什么必须有、以及为什么其他的坚决不加data底层的 NumPy 数组负责真正的数值存储。grad与data同形状的梯度数组初始化为全零。requires_grad布尔标记决定这个张量是否参与求导。叶子参数为真中间结果为假但中间结果仍然要记录运算否则链路断了。_backward一个无参函数负责把输出收到的梯度转成输入的梯度。_prev产生当前张量的父节点集合用于拓扑排序。注意这里没有shape、dtype、device这些属性。它们是data的派生属性写成property转发即可不需要存成独立字段。早期我给 Tensor 加了name字段方便调试结果某个环节的张量名被覆盖导致打印出来的计算图全是x反而更误导。调试信息应该用外部的字典维护不要污染数据结构本身。2.1 前向传播就是在悄悄构图很多教程把构建计算图和执行前向传播讲成两件事这个理解在动态图里是错的。动态图没有独立的构图阶段每一次运算都是一边算数一边把节点连起来。你写y x w b实际上做了三次构图一次矩阵乘、一次广播加、一次赋值。具体到代码每个运算函数的结构高度一致以加法为例def __add__(self, other): other other if isinstance(other, Tensor) else Tensor(other) out Tensor(self.data other.data, requires_gradFalse) out._prev {self, other} def _backward(): # 加法对两个输入的偏导都是 1直接透传上游梯度 if self.requires_grad: self.grad unbroadcast(out.grad, self.data.shape) if other.requires_grad: other.grad unbroadcast(out.grad, other.data.shape) out._backward _backward return out这段代码里有两个细节值得停一下。第一是而不是。为什么必须累加因为一个张量可能在计算图里被使用多次比如y x * x里x出现了两次两条路径的梯度都要汇入x.grad。如果用赋值后写入的会覆盖先写入的你得到的梯度永远是错的而且错得很隐蔽——因为数值大小看着合理只是偏小。这是绝大多数手写实现翻车的第一个点我在第一次跑通时也中招了。第二是unbroadcast。这是全项目最容易被跳过、也最容易出问题的一步下一节专门拆开讲。2.2 拓扑排序不是可选项是正确性前提反向传播必须按从输出到输入的逆序执行而_prev记录的是一个集合集合本身没有顺序。如果直接递归调用父节点的_backward会出现什么情况考虑y (a b) * c其中a b这个中间结果被记为t。当你对y调用_backward时它会让t和c执行各自的_backward而t的_backward又会去更新a和b。问题在于如果t还有另一条更长的路径连接到输出那么t.grad在它自己的_backward被执行时可能还没累加完整导致传下去的梯度偏小。标准解法是拓扑排序def backward(self): topo, visited [], set() def build(node): if id(node) not in visited: visited.add(id(node)) for parent in node._prev: build(parent) topo.append(node) build(self) self.grad np.ones_like(self.data) # 标量损失的初始梯度为 1 for node in reversed(topo): node._backward()这里有个 Python 层面的坑用id(node)而不是node本身做集合元素。原因是如果没有给 Tensor 定义__hash__和__eq__默认的对象哈希可以用但一旦你为了调试方便重写了__eq__比如想做x y逐元素比较对象就变得不可哈希set()直接抛异常。我建议从一开始就用id()成本几乎为零省掉一类诡异报错。3. 广播梯度形状对不上就是静默错误unbroadcast这个函数值得单独用一节来讲因为它是手写自动求导里唯一一个不报错、但结果全错的环节。NumPy 的广播规则会悄悄把(4, 3)和(3,)相加成(4, 3)反向时上游传来的梯度是(4, 3)而被加数的形状是(3,)——形状不匹配但 NumPy 的会尝试广播多数情况下会直接抛ValueError你能看到报错。真正致命的是另一种情况(4, 1)和(4, 3)相加反向时对(4, 1)的梯度是(4, 3)形状上能广播不会报错但语义完全错了。正确的做法是先按右侧对齐补 1再对所有多出来的和被补 1 的轴求和def unbroadcast(grad, target_shape): # 补齐维数差异 while grad.ndim len(target_shape): grad grad.sum(axis0) # 对原本是 1 的维度求和并保维 for axis, size in enumerate(target_shape): if size 1 and grad.shape[axis] ! 1: grad grad.sum(axisaxis, keepdimsTrue) return grad.reshape(target_shape)这里的keepdimsTrue不能省。省略之后求和得到的形状会少一维后续reshape可能因为元素总数一致而碰巧成功但维度语义已经错了等到下一层矩阵乘法时才会以另一个错误的形式爆发排查成本翻倍。我第一次实现时为了简洁去掉了keepdims然后在三层网络里追了六个小时。3.1 矩阵乘法反向的两个转置别背公式的反向传播公式经常被写成dA dC B.T、dB A.T dC很多人背下来就完事但换个场景立刻懵。我建议用一个不依赖记忆的推导方法看形状。设A形状(m, k)B形状(k, n)C形状(m, n)上游梯度dC也是(m, n)。现在要得到dA它必须是(m, k)手头只有dC形状(m, n)和B形状(k, n)能让(m, n)和某个B的转置相乘得到(m, k)的只有B.T形状(n, k)。同理dB要(k, n)只能用A.T形状(k, m)乘dC形状(m, n)。这个形状驱动的推导方法在实现自定义算子时特别好用尤其是涉及批量维、多头维的时候。我在实现一个带 batch 的注意力算子时就是靠形状推出来的三次转置顺序比翻公式快得多。3.2 参数共享场景下的梯度累加验证如果你的网络里有共享参数比如同一个 Linear 层在时间步上复用必须验证梯度累加是否正确。验证方法很朴素把步数设为 1用数值梯度对比再设为 2看解析梯度是否大致是两步梯度之和。如果发现解析梯度是正确值的二分之一说明某处用了赋值而非累加如果是两倍说明某处重复累加了。这两个信号非常明确几乎能定位到具体哪一行。4. 数值梯度校验花半小时省三天我现在的习惯是任何新的算子写完之后先不接网络单独做一次数值梯度校验通过了再往下写。这个习惯是一次惨痛教训换来的——当时我在一个四层网络上调试损失不下降的问题改了学习率、换初始化、加 BatchNorm折腾两天最后发现是某个激活函数在负半轴的梯度符号写反了。如果一开始就做单算子校验这个问题三十秒就能发现。4.1 中心差分与步长的选择数值梯度的标准做法是中心差分def numeric_grad(f, x, h1e-5): grad np.zeros_like(x) it np.nditer(x, flags[multi_index]) for _ in it: idx it.multi_index old x[idx] x[idx] old h f_plus f(x) x[idx] old - h f_minus f(x) x[idx] old grad[idx] (f_plus - f_minus) / (2 * h) return grad关于h的取值常见的说法是越大越好以避免浮点误差越小越好以避免截断误差。实际经验是在 float64 下取 1e-5 到 1e-6 之间相对误差通常在 1e-7 到 1e-8 量级。如果你用 float32 存储h取 1e-5 时中心差分的截断误差已经接近浮点精度相对误差会恶化到 1e-3 左右这时候你不能简单地判断实现有 bug需要先换成 float64 再测一遍。我自己的校验脚本一律强制 float64测完再切回 float32 训练。判断标准我用这样一张表相对误差结论 1e-7基本正确可以放心使用1e-7 ~ 1e-5通常是浮点累积误差可接受1e-5 ~ 1e-3可疑检查 h 的取值和数据类型 1e-3大概率实现有 bug逐项排查4.2 校验时的三个假阳性陷阱陷阱一ReLU 的折点。ReLU 在 0 点不可导如果扰动后的输入刚好跨过零点数值梯度会给出一个看起来错了的结果。解决办法是构造测试输入时避开 0 附近或者对 ReLU 单独用max(x, 0)的解析结果对照跳过数值法。陷阱二批量归一化的统计量。如果算子里包含对 batch 维求均值的操作扰动单个元素会同时影响均值和方差前向f_plus和f_minus对应的统计量不同中心差分的假设就被打破了。这种情况需要把统计量固定在测试版本里单独验证反向公式。陷阱三随机性。如果算子内部有 dropout 或者随机采样两次前向的结果不同中心差分直接失效。测试时必须把随机种子固定或者关掉随机分支。5. 训练循环优化器的状态管理比公式更重要很多教程讲优化器就到一个公式为止SGD 是p - lr * gMomentum 是v beta * v gAdam 是一堆系数。但真正写实现时难点根本不在公式而在状态的生命周期管理。因为优化器要在多个训练步之间保持状态而这些状态的初始化时机、形状、与参数的对应关系一旦搞错你的模型会在第 100 步左右毫无征兆地开始发散。5.1 为什么要延迟初始化状态我最初的做法是在优化器构造函数里直接把v和m初始化为全零数组。问题是构造函数调用时往往还没有挂上参数——参数是在构建模型之后才注册给优化器的。于是v的形状是空的第一次更新时用p.grad的形状去一个空数组NumPy 广播直接炸掉。修改后的做法是在第一次step()时检查状态是否存在不存在就按当前梯度的形状懒初始化class Adam: def __init__(self, params, lr1e-3, betas(0.9, 0.999), eps1e-8): self.params params self.lr, self.eps lr, eps self.b1, self.b2 betas self.state {} # id(p) - {m:..., v:..., t:...} def step(self): for p in self.params: g p.grad if g is None: continue s self.state.setdefault( id(p), {m: np.zeros_like(g), v: np.zeros_like(g), t: 0} ) s[t] 1 s[m] self.b1 * s[m] (1 - self.b1) * g s[v] self.b2 * s[v] (1 - self.b2) * (g * g) m_hat s[m] / (1 - self.b1 ** s[t]) v_hat s[v] / (1 - self.b2 ** s[t]) p.data - self.lr * m_hat / (np.sqrt(v_hat) self.eps)用id(p)做键而不是p本身原因和前面一样Tensor 重写了某些魔术方法之后可能不可哈希。用id还有一个好处是明确表达这是身份绑定不是值相等避免了两个数值相同的参数被误认为同一个。5.2 偏差修正这个细节值不值得写Adam 论文里的偏差修正m_hat m / (1 - beta1^t)我在第一版实现里直接跳过了理由是训练几千步之后 1 - 0.9^t 已经接近 1影响可以忽略。实测下来这个理由只对了一半对于 5000 步以上的长训练确实几乎无感但对于小规模快速验证比如我要在 200 步内确认模型能不能过拟合单个 batch偏差修正缺失会让前 20 步的有效学习率小得离谱损失曲线前段几乎是平的很容易被误判成模型没学到东西。我现在的做法是保留偏差修正但把它作为可开关选项。写完整版花的代码不超过五行省下的调试时间远超这个成本。5.3 从损失曲线读出实现问题损失曲线的形状对实现 bug 有很强的诊断价值我总结了几种典型形态曲线形态大概率原因完全水平不动学习率为 0、梯度全为 0、或用了赋值而非累加导致梯度被覆盖前几步正常随后突然飙高学习率过大、缺少梯度裁剪、或误差累计到发散剧烈震荡锯齿状batch 太小、学习率偏大、或者某层没做归一化单调下降但收敛值远高于预期容量不足、初始化方差太小、或损失函数实现有误下降后趋于平缓又缓慢上升过拟合或正则项系数设置不当其中完全水平不动是我最常见的问题几乎每次都指向某个梯度被覆盖掉了。6. 用单 batch 过拟合验收整条链路写完框架之后怎么确认它是真的能训练而不是看起来在跑我的标准验收动作是取一个固定的、只有一个 batch 的小数据集把模型训练到几乎零损失。这个动作来自一个朴素的判断——如果你的实现连单个 batch 都过拟合不了那它一定有问题跟数据规模、正则化、学习率调度都没有关系。具体做法是随机生成 32 条样本输入维度 20隐藏层 64输出 2 类交叉熵总参数量不到 2000。正常情况下学习率 1e-2 的 Adam 在 200 步内应该能把损失压到 1e-3 以下。如果压不下去排查顺序是固定的先手工验证一次前向的损失值用独立的 NumPy 代码算一遍交叉熵和框架输出对比再做单算子数值梯度校验最后检查优化器是否真的改了参数。6.1 参数真的被更新了吗这个检查听起来很蠢但我确实遇到过参数没更新的情况。检查方法极其简单训练前后各存一份参数的拷贝算一下两者差值的 L2 范数。before [p.data.copy() for p in model.params()] model.train_step(...) after [p.data.copy() for p in model.params()] print([np.linalg.norm(a - b) for a, b in zip(after, before)])如果打印出来是一串 0说明更新链路断了。常见原因有三种requires_grad标记没设对导致参数梯度根本没算优化器传进去的参数列表是空列表模型里params()方法返回的是子层参数的引用如果子层注册方式有问题会返回空或者step()里写了p.data p.data - lr * p.grad但p是重新创建的对象。最后这个坑很隐蔽因为赋值写法看起来完全正常实际是给局部变量重新绑定了没有影响原数组。6.2 扩规模时的内存与重复分配问题单 batch 跑通之后把 batch 放大到 256、层数加到 6 层很快就会遇到内存问题。原因是动态图默认会保留所有中间张量因为反向传播要用。参数只有几万个但中间激活值可能上百万个浮点数。我遇到的第一个内存问题不是爆炸而是训练变慢每步都要构造全新的计算图对象Python 对象创建和垃圾回收的开销在小模型上不明显但到了每步几千个中间节点的时候纯 Python 开销就超过了矩阵计算的耗时。我的处理方式很粗暴但有效把计算图节点上的_backward从闭包改成绑定到对象的方法减少闭包捕获变量的内存占用并且在不需要梯度的时候比如评估阶段用标记显式跳过_prev的记录省掉构图开销。实测在一个五层网络上评估阶段能省掉约 30% 的时间。这里还有个经验不要过早优化训练环路。我在第一周花了大量时间研究怎么避免中间张量的分配后来发现torch之所以快靠的是底层内存池和计算图优化器不是靠 Python 层面少建几个对象。在这个规模上正确的做法是先保证正确性接受比成熟框架慢十倍的现实然后去优化真正占时间的那几个矩阵乘法比如把(N, C)的逐样本循环换成一次矩阵乘。7. 我从这套实现里得到的几个具体结论做完这一遍之后有几个认知发生了变化写下来供参考。第一框架最难的从来不是数学公式。链式法则在纸上两行就推完了但在代码里你要处理形状对齐、重复使用导致的梯度累加、不可导点的分支、以及浮点精度带来的误判。这些东西在论文里一个都不会提但占了实现工作量的八成。第二累加语义是自动求导的生命线。我现在看任何手写的自动求导实现第一眼就去找grad 找到赋值号我就知道这里一定有问题。这个直觉在后来帮我看别人的实现时省了很多时间。第三验证手段的价值远高于实现手段。数值梯度校验这三十行代码我写了删、删了改前后优化了五六版它带来的价值超过项目里任何一段核心代码。它把我觉得应该是对的变成我确认它是对的。第四零依赖的目标本身是手段而不是目的。DeepZero 最终并不是要替代任何框架它的意义是把抽象的门槛拆掉一层。等你亲手把梯度从输出一路传回输入再看那些自动求导的文档感受是完全不同的——你不再需要去猜这里到底发生了什么因为那个过程你在自己的代码里已经看过几十遍了。