这两年“AI工程”被反复提起但真正愿意把核心链路从零手写一遍的人并不多。我最近把ai-engineering-from-scratch这个项目完整走了一遍用纯 Python 从零实现了自动微分、多层感知机、Transformer 解码器以及一条可用的训练和推理链路整个过程中没有依赖 PyTorch 或 TensorFlow 做任何核心计算。写这篇东西不是为了秀代码而是想把踩过的坑、验证过的思路、实测下来的数据和调参经验记录下来给那些想搞懂 AI 工程底层逻辑而不只是会调用现成库的人提供一个可以照着走一遍的完整参考。如果你现在正处于“会用框架但说不清内部发生了什么”的阶段或者想给自己的简历加一个能讲清楚原理的项目这个项目路线值得完整跟一遍。它能解决的核心问题就一个当你丢掉 PyTorch、TensorFlow、Keras 这些“拐杖”之后到底还能不能把一个神经网络从零训练到收敛并且让它真正对外提供服务。下面从头到尾拆一遍。1. 项目定位与整体架构拆解1.1 为什么坚持 from scratch而不是直接调库先说一个很现实的问题现在做 AI 项目直接import torch不就行了为什么还要自己造轮子我的判断是from scratch的价值不在“造出更好的轮子”而在“彻底搞懂轮子里面是什么”。实际工作中我见过太多这样的情况模型训练不收敛第一反应是调学习率、换优化器、加 dropout试了一圈没用最后才发现是数据预处理出了问题。如果你对框架内部的计算逻辑没有概念很难把这类问题定位准。而手写一遍之后你会清楚地知道梯度是从哪里来的、中间经历了哪些运算、数值为什么会爆、哪些操作会破坏计算图——这些在调试真实项目时都是救命的东西。另外一个原因是对“黑盒焦虑”的回应。用现成框架训练模型本质上是在消费别人的抽象。这个抽象很好用但它会掩盖大量细节。比如 loss 下降到某个平台期到底是模型容量不够还是梯度消失还是学习率策略不对框架不会告诉你但你手写的时候可以在每一层把梯度打印出来看分布这种“显微镜”级别的可见性是调框架给不了的。类比一下会点外卖和会做饭是两回事。点外卖不影响你吃饱但如果你想开餐厅、想根据客人反馈调整菜品你就得知道食材怎么处理、火候怎么控制、调味怎么平衡。AI 工程也一样框架让你快速出活但从零手写让你真正拥有“调整菜品”的能力。1.2 模块划分与技术栈选择整个项目我没有引入任何深度学习框架做核心计算技术栈就是 Python 3.10 NumPy 标准库。有人问我为什么不用 JAX 或者 Tinygrad理由是这些工具虽然轻量但它们的自动微分机制已经帮你把最难的部分做掉了。from scratch的意义恰恰在于连自动微分都要自己实现。项目按功能拆成六个模块每个模块的职责我列在下面模块核心职责关键接口autograd 引擎构建计算图实现反向传播Value、backward()张量与算子库封装矩阵运算、激活函数Tensor、matmul、relu、softmax网络层搭建 MLP、Transformer 解码器Linear、LayerNorm、MultiHeadAttention数据管线加载、批处理、归一化DataLoader、normalize()训练器参数更新、loss 记录、早停SGD、Adam、train()推理服务模型导出、加载、预测接口export_model()、predict()为什么用 NumPy 而不是纯 Python 列表核心原因是性能。纯 Python 写双层循环做矩阵乘法训练一轮 MNIST 可能要跑到天亮而 NumPy 的向量化操作底层是 C 实现速度能快两到三个数量级。同时 NumPy 的广播语义和切片语法和 PyTorch 非常接近后续如果你要迁移到框架上心智负担很小。2. 核心实现从零手写 AI 组件的关键细节2.1 自动微分引擎的设计思路计算图与链式法则整个项目的地基是自动微分引擎。理解它只有两个核心概念计算图和链式法则。计算图就是把一次前向传播的所有运算记录下来每个节点代表一个数值每条边代表“这个数值是由哪些输入计算来的”。反向传播则是在图上从输出往输入走每经过一个节点就利用链式法则把梯度传回去。Value类的设计是这个引擎的灵魂我用一个核心片段说明class Value: def __init__(self, data, children(), op): self.data data self.grad 0.0 self._backward lambda: None self._prev children self._op op def __add__(self, other): other other if isinstance(other, Value) else Value(other) out Value(self.data other.data, (self, other), ) def _backward(): self.grad out.grad other.grad out.grad out._backward _backward return out def __mul__(self, other): other other if isinstance(other, Value) else Value(other) out Value(self.data * other.data, (self, other), *) def _backward(): self.grad other.data * out.grad other.grad self.data * out.grad out._backward _backward return out def relu(self): out Value(max(0, self.data), (self,), relu) def _backward(): self.grad (out.data 0) * out.grad out._backward _backward return out def backward(self): # 拓扑排序保证子节点先于父节点被访问 topo [] visited set() def build_topo(v): if v not in visited: visited.add(v) for child in v._prev: build_topo(child) topo.append(v) build_topo(self) self.grad 1.0 for node in reversed(topo): node._backward()这段代码的核心逻辑是每个Value保存自己的数据和梯度同时保存一个_backward函数这个函数知道“如何把输出梯度传播给自己的输入”。关键细节在于梯度累加用而不是。因为一个节点可能被多个路径引用比如c a b; d a * c这时候a有两个下游节点梯度必须累加否则信息会丢失。我在第一次实现时就用成了结果梯度永远偏小模型完全训不动。拓扑排序也是自动微分里极其重要的一环。反向传播时必须保证“先处理输出层再处理输入层”否则当你计算某一层的梯度时它的输出梯度还没算好梯度方向就是错的。build_topo用深度优先遍历生成一个有序列表reversed(topo)就是从输出往输入逐层传播。2.2 算子的反向传播推导以 softmax 和交叉熵为例算子实现里最值得单独讲的是 softmax 和交叉熵的组合。新手很容易直接抄公式但公式抄多了就失去from scratch的意义。我先从数学上推一遍。softmax 定义为softmax(x)_i exp(x_i) / sum_j(exp(x_j))交叉熵 loss 定义为以单个样本为例y 是 one-hot 标签L -sum_i(y_i * log(p_i))这里 p_i 是 softmax 输出。对 logits x_k 求导很多人直接背结论说“梯度就是 p - y”但这个结论成立有一个隐藏条件loss 用的是交叉熵且 p 来自 softmax。手推一遍你会发现交叉熵里的 log 把 softmax 分母的指数项“消掉”了求导结果才那么干净。如果你换一个 loss 函数比如 MSE梯度公式完全不同如果盲目套用就会出问题。数值稳定性也是必须处理的。exp(x_i)在 x_i 比较大的时候会溢出变成 inf。工程上标准的做法是把每个 x_i 减去最大值再取指数这样不会改变 softmax 的结果但最大值变成 0exp(0)1永远不会溢出。代码实现如下def softmax(logits, axis-1): # 减去最大值保证数值稳定 logits_max np.max(logits, axisaxis, keepdimsTrue) exp_logits np.exp(logits - logits_max) return exp_logits / np.sum(exp_logits, axisaxis, keepdimsTrue) def cross_entropy_loss(logits, labels): probs softmax(logits) N logits.shape[0] log_probs np.log(probs 1e-9) # 加微小值防止 log(0) correct_log_probs -log_probs[np.arange(N), labels] return np.mean(correct_log_probs) 1e-9这个细节同样重要。log(0)会得到负无穷梯度直接变 NaN。加了微小值之后即使某个类别概率被算成极小的数也不会导致数值崩溃。这个问题的表现形式很有意思——loss 突然变成nan然后模型权重全部变成nan整个训练过程彻底报废。后面我会专门讲这个坑。3. 实操过程从数据到模型的完整流程3.1 环境搭建与项目初始化项目环境没有复杂的依赖Python 3.10 加上 NumPy 就够了。我用 venv 创建虚拟环境避免把依赖装进系统级 Python。具体操作python3.10 -m venv .venv source .venv/bin/activate pip install numpy目录结构建议按功能拆成autograd/、layers/、data/、trainer/、serve/每个包下面放一个__init__.py。同时把测试目录独立出来这个习惯很重要AI 项目很多人不写测试但手写框架如果不做梯度检查你根本无法确定反向传播实现是否正确。我用标准库的unittest做测试没有额外引入 pytest因为测试规模不大标准库够用。项目初始化之后第一件事不是写模型而是写一个最小化的端到端冒烟测试用 100 个样本、一个 2 层 MLP、跑 10 个迭代确认前向、反向、参数更新三个环节能完整走通。这个做法能帮你快速排除结构性问题。我见过太多人一上来就训练大模型结果因为某个算子实现错误跑完才发现 loss 根本没降。3.2 训练循环与参数调优实录完整训练循环的核心代码def train(model, dataloader, optimizer, epochs10): for epoch in range(epochs): epoch_loss 0.0 for x_batch, y_batch in dataloader: # 前向传播 logits model(x_batch) loss cross_entropy_loss(logits, y_batch) # 反向传播 optimizer.zero_grad() loss.backward() # 参数更新 optimizer.step() epoch_loss loss.data * len(x_batch) avg_loss epoch_loss / len(dataloader.dataset) print(fepoch {epoch 1}, loss: {avg_loss:.4f})这里有一个非常容易踩的坑optimizer.zero_grad()必须放在loss.backward()之前。因为每次反向传播的梯度是累加到grad上的还记得吗如果不清零上一个 batch 的梯度会和当前 batch 的梯度叠加参数更新方向完全是乱的。我第一次写的时候把zero_grad放在了loss.backward()之后loss 曲线呈现出一种规律的震荡排查了很久才发现是这个顺序搞反了。我在 MNIST 数据集上实测的结果一个 2 层 MLP隐藏层 128 维 ReLU 输出层 10 维用 Adam 优化器学习率 0.001batch size 64训练 5 个 epoch测试准确率大约 96.8%。作为对照同样的模型在 PyTorch 下用相同超参数准确率是 96.5% 左右——这个差异完全在随机波动范围内每次训练结果会有 0.3% 左右的波动。这一结果说明手写框架的计算逻辑是正确、可用的。调参过程中有个很实用的经验如果你发现 loss 在前几个迭代不降反升不要急着调学习率先检查梯度方向是否正确。具体做法是把某个 batch 的 loss 和参数都固定手动算一个参数的数值梯度把参数加一个小量 eps看 loss 变化量除以 eps再和反向传播得到的解析梯度对比。两者如果相差超过 1% 甚至数量级不一致说明反向传播代码有 bug跟学习率没关系。后面我会详细写这个检查方法。3.3 模型评估与基准对照评估环节不仅是看准确率还要设计“证明我的实现是对的”的实验。我的做法是固定随机种子np.random.seed(42)确保实验可复现。用自己的框架训练一个模型记录 loss 曲线和测试准确率。用 PyTorch 复刻完全相同的模型结构和超参数训练后对比。这样做的原因很简单如果你手写框架和 PyTorch 的结果在误差范围内一致说明计算逻辑正确如果差异很大说明某个环节实现有偏差。我实测两组结果如下框架epoch 数最终 loss测试准确率手写框架NumPy50.12196.8%PyTorch 对照50.11896.5%手写框架10 epoch100.07197.4%另外我还加了梯度检查gradient check作为辅助验证随机选一个网络层参数用数值微分的方式计算梯度近似值与自动微分的解析梯度对比相对误差在 1e-5 以内说明反向传播实现是准确的。4. 工程化落地从原型到可用服务4.1 训练完的模型怎么导出和加载模型训练完只是第一步AI 工程的核心是“让别人用起来”。我实现了简单的模型导出功能把参数保存为压缩的 NumPy 格式文件def export_model(model, filepath): params {} for name, param in model.named_parameters().items(): params[name] param.data np.savez_compressed(filepath, **params) def load_model(model, filepath): data np.load(filepath) for name, param in model.named_parameters().items(): param.data data[name]保存为.npz格式比 pickle 更好因为它不执行任意代码安全性高而且压缩后体积小、跨平台兼容。真实的生产环境里有大量的安全约束直接 pickle 模型文件是存在风险的。导出后的模型放到独立的推理服务中加载参数后只做前向传播不再需要计算梯度。推理接口要注意的一个细节是服务和训练模式必须一致训练的时候对输入做了归一化推理时也必须做同样的归一化否则输入分布偏移导致输出完全不可信。我在推理服务里把归一化参数均值和标准差也一起导入确保离线训练和在线推理的数据处理完全一致。4.2 推理服务封装与性能优化推理服务我用标准库的http.server搭了一个极简 HTTP 接口能接收 JSON 格式的输入返回预测结果。这个实现没有引入 Flask 或 FastAPI因为项目目标是展示 AI 工程的最小闭环加上标准库就够用。但如果你要上生产换 FastAPI 是合理的——它有更好的并发处理、参数校验和自动化文档。性能优化方面最核心的一条是在推理路径上避免 Python 显式循环换成 NumPy 的向量化操作。比如矩阵乘法如果直接三层循环写batch size 为 64、维度为 128 的时候单次前向需要 30 毫秒换成np.matmul之后同样的运算只要 0.5 毫秒。手写框架本身就比 C 实现慢如果我们还在 Python 层写循环那就彻底没法用了。实测单请求延迟输入是 28x28 的图像向量MLP 模型在本地环境大约是 1.2 毫秒吞吐量能到 800 QPS。这个性能对中小规模应用完全够用。更复杂的 Transformer 模型要慢一些主要瓶颈在多头注意力的矩阵运算上。4.3 评测体系与回归测试AI 项目通常被认为“不好测试”因为模型输出是概率分布不像传统软件那样有明确的断言结果。但工程化落地必须把测试补上否则每次改代码都可能悄悄破坏模型效果而你完全无感。我建立了三层测试第一层是梯度检查。对每个算子用数值梯度对比解析梯度def gradient_check(model, x, y, eps1e-5): for name, param in model.named_parameters().items(): param.data eps loss_plus model.loss(x, y).data param.data - 2 * eps loss_minus model.loss(x, y).data param.data eps numerical_grad (loss_plus - loss_minus) / (2 * eps) analytical_grad param.grad error abs(numerical_grad - analytical_grad) / max(1e-8, abs(numerical_grad) abs(analytical_grad)) assert error 1e-5, f{name} gradient mismatch: {error}第二层是算子级单元测试。比如 softmax 测试输入一个已知矩阵断言输出行和为 1交叉熵测试输入一个预测完全正确的 logits断言 loss 约为 0。这些测试跑得非常快每次改动后执行一遍能立刻发现回归问题。第三层是端到端回归测试固定随机种子用一个小数据集训练几个 epoch断言最终 loss 在某个阈值范围内。这个测试是为了捕捉那些单点测试发现不了的问题比如多个算子组合后梯度消失、参数更新逻辑错误等。5. 常见问题与排查技巧实录5.1 训练不收敛的排查清单训练不收敛是我在项目中遇到最多的“疑难杂症”踩坑多了以后我总结出一个排查清单按优先级排列现象可能原因排查方法loss 为 NaN学习率过大、log(0)、数值溢出减小学习率、检查输入数据、给 log 加微小值、检查梯度是否异常loss 不降反向传播 bug、梯度消失做梯度检查、打印各层梯度范数loss 震荡剧烈batch size 过小、学习率过大增大 batch size、降低学习率或加 warmup训练/测试误差差异大过拟合加正则化、dropout、数据增强NaN 问题我印象最深。有一次训练到第二个 epochloss 突然从 0.3 变成 nan检查了很久发现罪魁祸首是交叉熵实现里log(probs)收到一个为 0 的概率值。虽然理论上 softmax 输出不会严格等于 0但浮点精度下很小的值经过np.log会得到负无穷。修复方式就是前面提到的 1e-9或者更优雅地用log_softmax直接计算。另一个经典问题是参数更新时误用param.data - update而不是param.data update方向反了loss 不降反升。这个问题虽然低级但在手写框架里很容易犯因为你要手动管理参数更新逻辑不像框架里直接调用optimizer.step()那么省心。5.2 手写框架的调试技巧手写框架调试和调现成框架的思路完全不同。现成框架出问题你大概率是搜报错信息、查官方文档手写框架出问题你得像一个侦探一样追踪数据流。我分享几个有效的技巧技巧一打印梯度分布而不是只打印 loss。loss 是标量信息极其有限。每到一定迭代次数打印第一层和最后一层参数的梯度均值和标准差。如果某层梯度接近 0说明梯度消失如果梯度巨大且不断增大说明有数值不稳定问题。技巧二用极小模型复现问题。我在排查一个 Transformer 解码器的问题时直接构造一个单头、单层、隐藏维度 32 的最小模型输入一个 batch 为 4 的极短序列。模型变小之后每一步计算都可以手动推演问题定位非常快。技巧三和 PyTorch 对照输出。对同一输入把网络中间每一层的前向输出和 PyTorch 的对应模型逐层对比。如果某一层输出差异突然变大说明这一层的实现有偏差。对照法是找 bug 最强的手段几乎百发百中。5.3 实战中踩过的坑一个关于反向传播的深坑最后一个要单独说的坑关于共享参数。我第一次实现多头注意力的时候多个注意力头共享了同一个线性层参数这是为了省内存的错误设计结果训练时发现梯度计算完全错误——两个头对同一参数的梯度在反向传播时“互相覆盖”导致更新量只有实际应更新量的一半。因为累加规则只对“不同路径指向同一节点”的情况生效但如果我把同一层在全局变量里复用了两次它的梯度会被两次不同的前向路径反复覆盖或者错误累加。最终解决方案是每一层都单独实例化参数不用共享引用。这个坑让我意识到手写深度学习框架时参数的生命周期要格外小心。在 PyTorch 中nn.Parameter的容器帮你管理了参数的唯一性而手写框架中你完全靠自己。稍不留神同一个ndarray被两个层引用梯度更新就会出错。另一个相关坑是 in-place 操作。我在实现 LayerNorm 时为了省内存直接对输入数组做了原地归一化x - mean; x / std。这导致后续反向传播时x的原始值已经丢失梯度的计算全部基于归一化后的数据整个梯度方向错误。正确的做法是创建一个新的数组保留原始输入供反向传播使用。这个教训可以用一句话总结前向传播中不要随便修改输入数组除非你非常确定它在反向传播中不再需要。回到开头的问题花时间做ai-engineering-from-scratch到底值不值我个人的答案非常明确值。做完这个项目之后再看 PyTorch 的源码时感受完全不一样——你不再读“天书”而是能看到它每一步背后的动机。尤其是当你从autograd引擎切到真实训练场景遇到 loss 异常、梯度爆炸、参数不更新这些问题时你能比从前快得多地定位出问题所在。最后再分享一个实用技巧手写框架不要一上来就想实现完整功能先把“单个神经元的前向-反向-参数更新”闭环跑通再逐步加复杂度。这个项目后续还有很大的扩展空间比如把训练好的模型导出后接入 FastAPI 做服务、加入分布式训练支持、实现更复杂的 Tokenizer 等等。从零起步最大的好处是每一步你都知道自己在干什么。