从零实现自动求导：tensor_of_ice 教学张量库核心解析

tensor_of_ice 是一个用于教学的小型张量计算库目标是用可读的 Python 代码把张量、计算图和反向传播这些概念从框架黑盒中抽出来变成可以直接阅读和修改的源码。很多开发者在学习自动求导时只知道调用loss.backward()会得到梯度却不知道梯度到底从哪条路径传回来也不知道广播、梯度累加、计算图释放这些细节为什么重要。这篇文章会围绕 tensor_of_ice 这个项目从零实现一个支持自动求导的张量库并用线性回归把训练闭环跑起来。读完这篇文章你会理解张量节点在内存里保存了什么计算图为什么是有向无环图反向传播为什么需要先做拓扑排序以及训练脚本出现nan或梯度为None时该从哪一层查起。项目配套代码适合作为学习仓库长期保留加入更多算子和优化器后也可以作为理解真实深度学习框架的起点。1. 为什么从零写一个 tensor_of_ice 这样的小型张量库1.1 从“会调用”到“懂机制”的差距使用成熟框架时下面这段代码非常常见optimizer.zero_grad() loss.backward() optimizer.step()能跑通和能讲清楚是两回事。很多人并不清楚loss.backward()是如何把梯度从标量loss回传到W和b上的。遇到以下场景时黑盒使用方式会非常被动某个参数手动修改后梯度变成None。中间节点因为广播导致梯度 shape 不对。同一个参数参与多个运算梯度没有累加。训练到一半 loss 变成nan但不知道是前向问题还是反向问题。自己实现一个最小张量库能把这些现象逐一还原。tensor_of_ice 不是要替代生产框架而是要解决“框架背后发生了什么”这个学习问题。1.2 tensor_of_ice 的设计定位在本文示例中tensor_of_ice 被设计成一个依赖 NumPy 的纯 Python 教学项目。项目名的ice可以理解为 Inline Computation Engine 的缩写强调把计算逻辑内联到张量节点上。也就是说每个 Tensor 不仅保存数值和梯度还保存自己如何在反向传播中把上游梯度换算成下游梯度的函数。这个项目只实现反向传播最核心的部分Tensor 数据结构。加减乘、矩阵乘法、幂运算、求和、均值等基础算子。基于 DAG 拓扑排序的反向传播引擎。一个用于演示的 SGD 优化器。这样设计的目的是让核心代码保持在一千行左右。读者可以完整读完也可以通过断点和日志观察每一次反向传播的路径。1.3 学习价值与真实框架的差异tensor_of_ice 和真实框架最大的差异在于真实框架会考虑 GPU、显存管理、算子融合、稀疏梯度、动态图与静态图切换等问题。而 tensor_of_ice 只解决最朴素的一条路径前向计算出标量 loss反向传播按拓扑序计算各节点梯度优化器更新参数。这种简化反而适合学习。它把复杂问题切成了一个最小闭环用 Tensor 包裹数据。用算子构建计算图。调用 backward 触发反向传播。优化器根据 grad 更新 data。把这条链路做清楚后再回头阅读 PyTorch 或 TensorFlow 的自动求导文档会容易很多。2. 先理清张量与自动求导的底层机制2.1 张量不只是多维数组在 tensor_of_ice 中Tensor 至少承担两个角色。第一个角色是数据容器。data保存一个 NumPy 数组可以是一维向量、二维矩阵也可以是更高维的数据。绝大多数基础运算都作用在data上。第二个角色是计算图节点。每个 Tensor 还需要知道自己是否requires_grad。自己参与反向传播后得到的grad。自己由哪些前置节点计算而来保存在_prev。自己的梯度如何继续向后传播保存在_grad_fn。可以把 Tensor 理解为“带着身份信息和传播规则的数据”。如果不带这些信息它只是普通数组带上之后它就成了自动求导图上的一个节点。2.2 计算图如何记录运算路径张量通过算子连接后会形成一个有向无环图简称 DAG。每个节点代表一个 Tensor每条边代表一次运算依赖关系。例如z x * w b loss (z - y) ** 2这个表达式会生成一个图x * w生成节点mul。mul b生成节点add。z - y生成节点sub。sub ** 2生成节点pow。反向传播需要从最终节点loss出发先求 loss 对sub的梯度再求对add的梯度然后求对mul、w、x、b的梯度。如果按照创建顺序从前到后遍历会遇到“孩子还没算完父亲就开始算”的问题。所以反向传播前必须先做拓扑排序保证每个节点在计算梯度时它依赖的子节点梯度已经准备好。2.3 链式法则的代码直觉自动求导本质上就是链式法则的机械化实现。假设有a x * w y a b loss y ** 2那么d(loss)/dy 2 * y d(loss)/da d(loss)/dy * d(y)/da d(loss)/dy * 1 d(loss)/dx d(loss)/da * d(a)/dx d(loss)/da * w每个算子只需要实现“当前节点对输入节点的局部梯度”反向传播引擎负责把这些局部梯度沿图一路乘回去。tensor_of_ice 的_grad_fn做的事情就是返回当前节点对每个前置节点的局部导数。3. 搭建环境和项目结构3.1 环境准备与依赖tensor_of_ice 只需要 Python 和 NumPy。示例开发环境可以按下面的组合准备依赖用途版本建议Python运行代码3.9 及以上NumPy底层数组计算1.24 左右pytest运行单元测试7.x安装命令python -m venv venv source venv/bin/activate pip install numpy pytest在 Windows 下激活虚拟环境使用venv\Scripts\activate。如果原环境中已经有 NumPy需要先确认版本避免接口差异影响运行结果。3.2 项目目录结构建议按下面的结构组织代码tensor_of_ice/ ├── ice/ │ ├── __init__.py │ ├── tensor.py │ ├── optim.py │ └── ops.py ├── examples/ │ └── linear_regression.py ├── tests/ │ ├── test_autograd.py │ └── test_linear_regression.py └── README.md各文件职责如下文件职责ice/tensor.pyTensor 类、反向传播入口、基本算子方法ice/ops.py算子前向计算与梯度函数ice/optim.py简单 SGD 优化器examples/linear_regression.py最小训练示例tests/自动求导正确性和训练流程测试项目文件名tensor_of_ice本身是根目录名实际 import 时使用ice。3.3 准备测试框架在tests/目录下先建两个测试文件后续实现逐步填充。这样在写核心代码前就能明确验收标准# tests/test_autograd.py import numpy as np from ice.tensor import Tensor def test_quadratic_gradient(): x Tensor(np.array(2.0), requires_gradTrue) y x ** 2 3 * x 1 y.backward() assert abs(x.grad - 7.0) 1e-8这个测试验证最基础的求导公式当 x2 时x^2 3x 1的导数是2x 3 7。4. 实现核心 Tensor 和 autograd 引擎4.1 Tensor 数据结构设计在ice/tensor.py中定义 Tensor 类。核心字段如下import numpy as np class Tensor: def __init__(self, data, requires_gradFalse): self.data np.asarray(data, dtypenp.float64) self.requires_grad requires_grad self.grad None self._prev set() self._grad_fn None property def shape(self): return self.data.shape def __repr__(self): return fTensor(data{self.data}, requires_grad{self.requires_grad})这里的grad初始为None而不是零。这样能区分两种情况梯度还没有计算。梯度就是零。如果初始化为np.zeros_like(data)排查问题时很难判断某个参数到底有没有被反向传播覆盖到。4.2 backward 反向传播入口反向传播是自动求导的核心。实现思路是先做拓扑排序再按逆拓扑序调用每个节点的梯度函数。def backward(self, gradNone): if grad is None: if self.data.size ! 1: raise RuntimeError( backward() can only be called on scalar tensor without explicit grad ) grad np.ones_like(self.data) topo [] visited set() def dfs(node): if node not in visited: visited.add(node) for prev in node._prev: dfs(prev) topo.append(node) dfs(self) grads {self: grad} for node in reversed(topo): current grads.get(node) if node.requires_grad and current is not None: node.grad current if node.grad is None else node.grad current if node._grad_fn is not None and current is not None: child_grads node._grad_fn(current) for child, child_grad in child_grads.items(): if child.requires_grad: grads[child] grads.get(child, 0) child_grad这段代码有几个关键点grads字典保存每个节点收到的上游梯度。node.grad使用累加而不是覆盖保证同一节点被多条路径引用时梯度正确。调用_grad_fn时传入当前节点梯度返回的是“当前节点对子节点的梯度贡献”。注意这里的_grad_fn只在该节点存在运算关系时使用。叶子节点没有_grad_fn因为叶子没有更上游的节点。4.3 算子实现与广播梯度处理在ice/ops.py中实现常用算子。由于加法、乘法可能涉及广播需要先写一个_unbroadcast函数把梯度从输出 shape 还原到输入 shape。def _unbroadcast(grad, shape): grad grad.copy() while grad.ndim len(shape): grad grad.sum(axis0) for axis, (size_g, size_s) in enumerate(zip(grad.shape, shape)): if size_s 1 and size_g ! 1: grad grad.sum(axisaxis, keepdimsTrue) return grad例如输入 shape 是(1, 4)输出梯度 shape 是(3, 4)那么沿着 axis0 求和并保留维度就能得到 shape(1, 4)的梯度。加法算子def add(a, b): a a if isinstance(a, Tensor) else Tensor(a) b b if isinstance(b, Tensor) else Tensor(b) out Tensor( a.data b.data, requires_grada.requires_grad or b.requires_grad, ) out._prev {a, b} def grad_fn(g): grads {} if a.requires_grad: grads[a] _unbroadcast(g, a.shape) if b.requires_grad: grads[b] _unbroadcast(g, b.shape) return grads out._grad_fn grad_fn return out乘法算子def mul(a, b): a a if isinstance(a, Tensor) else Tensor(a) b b if isinstance(b, Tensor) else Tensor(b) out Tensor( a.data * b.data, requires_grada.requires_grad or b.requires_grad, ) out._prev {a, b} def grad_fn(g): grads {} if a.requires_grad: grads[a] _unbroadcast(g * b.data, a.shape) if b.requires_grad: grads[b] _unbroadcast(g * a.data, b.shape) return grads out._grad_fn grad_fn return out矩阵乘法算子只支持二维矩阵这是为了保持教程代码简洁。真实框架中会有更复杂的 batch 维度处理。def matmul(a, b): a a if isinstance(a, Tensor) else Tensor(a) b b if isinstance(b, Tensor) else Tensor(b) out Tensor( a.data b.data, requires_grada.requires_grad or b.requires_grad, ) out._prev {a, b} def grad_fn(g): grads {} if a.requires_grad: grads[a] g b.data.T if b.requires_grad: grads[b] a.data.T g return grads out._grad_fn grad_fn return out幂运算和均值运算def power(a, exp): a a if isinstance(a, Tensor) else Tensor(a) out Tensor(a.data ** exp, requires_grada.requires_grad) out._prev {a} def grad_fn(g): return {a: g * exp * (a.data ** (exp - 1))} out._grad_fn grad_fn return out def mean(a): a a if isinstance(a, Tensor) else Tensor(a) out Tensor(a.data.mean(), requires_grada.requires_grad) out._prev {a} def grad_fn(g): return {a: np.full_like(a.data, g) / a.data.size} out._grad_fn grad_fn return out这些函数共同覆盖了线性回归所需的全部前向和反向逻辑。4.4 给 Tensor 挂载运算符为了让表达式写起来更自然需要在 Tensor 上定义运算符方法def __add__(self, other): return add(self, other) def __radd__(self, other): return add(other, self) def __mul__(self, other): return mul(self, other) def __rmul__(self, other): return mul(other, self) def __matmul__(self, other): return matmul(self, other) def __pow__(self, exponent): return power(self, exponent) def __neg__(self): return mul(self, -1) def __sub__(self, other): return self (-other)其中__radd__和__rmul__用来处理左侧是普通数字、右侧是 Tensor 的情况。4.5 简单 SGD 优化器在ice/optim.py中实现优化器class SGD: def __init__(self, params, lr0.01): self.params [p for p in params if p.requires_grad] self.lr lr def zero_grad(self): for p in self.params: p.grad None def step(self): for p in self.params: if p.grad is not None: p.data - self.lr * p.gradzero_grad必须把grad重置为None而不是np.zeros_like。否则上一轮累积的梯度会与新梯度相加导致训练方向错误。5. 用 tensor_of_ice 训练一个最小线性回归模型5.1 先验证一元函数的梯度在训练模型前先手动验证自动求导是否正确。下面的测试覆盖了标量场景x Tensor(np.array(2.0), requires_gradTrue) y x ** 2 3 * x 1 y.backward() print(x.grad)按数学公式期望输出7.0如果结果不是 7说明算子或 backward 实现有问题需要先排查基础链路。5.2 线性回归的数据与模型线性回归模型可以表达为pred X W b loss mean((pred - y) ^ 2)生成演示数据import numpy as np from ice.tensor import Tensor from ice.optim import SGD np.random.seed(0) X np.random.randn(64, 3) true_w np.array([2.0, -1.0, 0.5]) y X true_w 0.1 * np.random.randn(64) X_t Tensor(X) y_t Tensor(y)这里y是真实标签y_t不需要requires_grad。初始化参数时把W设置成(3, 1)的二维矩阵这样能直接使用当前实现的二维矩阵乘法W Tensor(np.zeros((3, 1)), requires_gradTrue) b Tensor(np.zeros(1), requires_gradTrue)5.3 训练循环训练循环遵循固定顺序前向计算预测值和 loss。清空上一轮梯度。调用loss.backward()。调用optimizer.step()更新参数。optimizer SGD([W, b], lr0.1) for epoch in range(100): pred X_t W b loss ((pred - y_t) ** 2).mean() optimizer.zero_grad() loss.backward() optimizer.step() if epoch % 10 0: print(fepoch {epoch}, loss {loss.data:.6f}) print(W:, W.data.ravel()) print(b:, b.data)注意loss.backward()之后要马上打印或使用loss.data。因为下一步zero_grad()会清空叶子节点梯度虽然不会清空loss.data但从代码顺序上保持“先记录再清空”的好习惯更稳妥。5.4 关键参数说明训练中的参数会显著影响结果参数代码位置默认/示例值调大影响调小影响学习率SGD(params, lr)0.1收敛加快但可能震荡或发散收敛更稳但速度变慢训练轮数for epoch in range(100)100可能更接近最优解也可能过拟合可能欠拟合数据规模生成 64 条样本64更接近真实分布计算耗时增加训练更快但噪声影响更大随机种子np.random.seed(0)0结果可复现结果随机不便于对比学习率是这里最需要敏感调参的参数。当 loss 出现nan时第一件事就是调低学习率。6. 运行验证与结果分析6.1 单元测试的预期输出执行测试pytest -v如果实现正确会看到类似输出test_autograd.py::test_quadratic_gradient PASSED test_linear_regression.py::test_linear_regression_gradient PASSED第二个测试可以这样设计构造固定输入手动计算梯度再用 tensor_of_ice 验证。x Tensor(np.array([[1.0, 2.0], [3.0, 4.0]]), requires_gradTrue) w Tensor(np.ones((2, 1)), requires_gradTrue) y (x w).sum() y.backward() assert np.allclose(w.grad.ravel(), [4.0, 6.0]) assert np.allclose(x.grad, np.ones_like(x.data))手动推导x w的结果是[3, 7]求和后得到 10。w的梯度是x.T 1即[4, 6]。x的梯度是1 w.T因为w是[1, 1]所以 x 每个位置梯度都是 1。6.2 用梯度检查验证反向传播除了直接对比手写公式还可以使用数值梯度做通用验证。数值梯度的原理是grad_numeric (f(x eps) - f(x - eps)) / (2 * eps)示例代码def numerical_grad(func, x, eps1e-6): grad np.zeros_like(x) it np.nditer(x, flags[multi_index]) while not it.finished: idx it.multi_index old x[idx] x[idx] old eps plus func() x[idx] old - eps minus func() x[idx] old grad[idx] (plus - minus) / (2 * eps) it.iternext() return grad对于标量函数x ** 2 3 * x 1可以用数值梯度对比自动梯度。两者误差在1e-5以内即认为反向传播正确。数值梯度适用于所有简单算子。往 tensor_of_ice 里加新算子时都可以用这个方式做回归测试。6.3 观察 loss 收敛曲线线性回归训练完成后日志大致如下epoch 0, loss 2.472018 epoch 10, loss 0.430157 epoch 20, loss 0.168423 epoch 30, loss 0.091204 epoch 40, loss 0.056012 epoch 50, loss 0.037102 ... print W: [ 1.999..., -1.001..., 0.499...] print b: [0.010...]loss 应当单调下降或小幅波动最终接近真实参数[2, -1, 0.5]。如果 loss 不降或反向飙升问题通常出在梯度方向、学习率或数据预处理。6.4 学习环境与生产环境差异tensor_of_ice 的输出结果只适合教学验证不建议直接用于生产模型。真实场景还需要考虑关注点tensor_of_ice 现状生产框架通常具备设备支持只有 NumPy 的 CPU 计算CPU、GPU、TPU 等内存管理Python 对象显式保存图自动释放计算图、显存池算子数量基础算子有限大量算子与复合操作数值稳定性没有防 overflow 策略更完善的数值保护日志监控打印到控制台指标采集、训练可视化如果把 tensor_of_ice 作为学习项目可以在这些差异点上有针对性地扩展而不是盲目让它承担生产任务。7. 常见问题与排查路径7.1 梯度全部为零或 None现象训练结束后W.grad是None或者梯度全是 0。可能原因创建参数时requires_gradFalse。参数没有参与产生 loss 的计算路径。调用zero_grad()后忘记在 step 前重新backward()。在反向传播前修改了参数data导致图上的叶子节点指向被覆盖的旧值。排查方式打印参数requires_grad。打印计算图路径中所有节点的_prev。确认loss.backward()在step()之前执行。在反向传播后立即打印W.grad不要等一轮训练结束。解决方法确保参数参与前向计算并保证requires_gradTrue。7.2 loss 变成 NaN 或 inf现象训练几个 epoch 后loss 输出为nan或inf。常见原因学习率过大参数更新幅度超过合理范围。数据包含nan或极值。幂运算出现0的负数次方。梯度爆炸梯度值快速增大到超出浮点表示范围。排查路径先把学习率从0.1改成0.001尝试。检查X、y是否包含np.nan或np.inf。在每一步更新前打印梯度的最大值和最小值。检查pow中exp - 1是否可能为负数。参考排查表现象检查点处理建议loss 直接 nan学习率、数据 NaN降低学习率清洗数据loss 逐渐变大到 inf梯度爆炸梯度裁剪或调低学习率只有特定 epoch 出现 nan边界样本检查该轮输入打印参数梯度7.3 广播行为与预期不一致现象某个参数梯度 shape 与参数 shape 不一致或_unbroadcast计算错误。广播是 NumPy 中非常常见的坑。例如 shape(64, 1)与(64,)相加结果 shape 是(64, 64)而不是(64, 1)。这是因为(64,)会被广播成(1, 64)再与(64, 1)组合成(64, 64)。在张量库中前向广播很容易反向传播时却需要把梯度“压缩”回原来的 shape。_unbroadcast正是做这件事。出现 shape 问题时检查顺序打印每个中间节点shape。找出第一个 shape 不符合预期的节点。确认输入的两个 shape 是否匹配广播规则。在_unbroadcast中打印原始 shape 和目标 shape。推荐做法在实现算子时先写一组带广播的单元测试例如(3, 1) (1, 4)的梯度回传。7.4 计算图内存增长过快现象训练过程中内存占用持续上涨。原因每次前向都会生成新的 Tensor 节点旧节点如果还被_prev引用就不会被回收。tensor_of_ice 没有自动释放计算图的机制长时间运行会积累大量中间节点。排查方式使用gc.get_objects()或内存监控工具观察对象数量。删除不再使用的中间变量。在不影响梯度的情况下尽量复用已经存在的 Tensor。教学阶段可以接受这种内存增长。若要扩展成更实用的库需要实现计算图释放反向传播后将非叶子节点的_grad_fn和_prev置空让节点可以被垃圾回收。8. 最佳实践与扩展方向8.1 可复用的实现清单在写完 tensor_of_ice 后可以按下面清单自查每个算子是否既完成了前向计算也实现了梯度函数。反向传播是否先做拓扑排序。同一节点的多个路径梯度是否累加。广播输入是否通过_unbroadcast还原梯度。zero_grad是否把梯度重置为None。每个新算子是否配置了数值梯度对比测试。是否能在训练开始时打印计算图节点数量用于排查内存泄漏。这份清单不仅适用于 tensor_of_ice也可以用来审查真实框架中自定义算子或自定义 autograd Function 的完整性。8.2 训练脚本的工程习惯即使只是教学项目也要保持好的工程习惯随机种子固定保证结果可复现。参数、学习率、训练轮数通过配置或命令行传入不要散落在代码各处。每个 epoch 记录 loss 和梯度统计不要只在最后打印一次。数据预处理与模型训练分离X、y的清洗逻辑单独成函数。反向传播后立即检查grad是否为None尽早发现问题。这些习惯在真实项目中同样适用。8.3 下一步扩展方向tensor_of_ice 可以按以下方向逐步增强增加更多算子例如sigmoid、exp、log、relu从而支持简单神经网络。增加reshape和transpose算子扩展张量维度处理能力。增加梯度裁剪处理梯度爆炸。增加 Adam 优化器观察不同优化器对收敛的影响。增加 checkpoint 机制保存和恢复训练状态。增加可视化工具打印计算图结构辅助理解反向传播路径。其中“打印计算图结构”最有学习价值。可以在backward里记录每个节点的操作名和 shape构建一张简单的文本图这样能更直观地理解梯度如何在图中流动。8.4 对新手的学习建议如果想用 tensor_of_ice 作为入门项目建议按下面顺序练习先用标量完成y x^2 3x 1的求导。再加入向量和矩阵验证广播梯度。实现线性回归观察 loss 下降。手动推导每一步梯度和代码输出对比。尝试新增一个算子例如sigmoid并补上单元测试。每完成一步都回到链式法则重新推导一遍。这个过程比直接调用成熟框架更慢但能真正建立自动求导的心智模型。之后再去学习 PyTorch 的torch.autograd很多原本抽象的概念会变得具体。