1. 为什么我要从零手搓一套AI工程流水线

第一次看到ai-engineering-from-scratch这个项目名的时候，我正被一堆“调包侠”式的教程搞得有点烦。满屏都是import torch、from transformers import ...，跑通一个 demo 只要十分钟，可真要把模型塞进生产环境，问题就全冒出来了：显存炸了、推理延迟飘了、batch 一改结果就对不上、日志里全是 NaN 却不知道从哪一层开始崩的。这些坑，光靠调库是填不平的。

ai-engineering-from-scratch这个标题，核心不在“AI”，而在“from scratch”——从零开始。它要解决的不是“怎么调用一个现成模型”，而是“怎么把 AI 这件事从数学公式一路搭到能跑、能测、能上线的工程系统”。这里面涉及的核心领域其实横跨了三块：深度学习底层原理、工程化实现、以及系统级性能优化。适合谁来参考？我的判断是三类人：一是刚学完理论、想动手把反向传播真正写一遍的学生；二是天天调库、但遇到性能瓶颈就束手无策的中级工程师；三是想搞清楚“框架到底帮我做了什么”的技术负责人。

我自己的背景是做了七八年后端和数据处理，转 AI 工程大概三年。这三年里我最大的体会是：你调库时省下的每一分钟，都会在排查线上问题时加倍还回来。所以当我决定认真啃一遍从零实现的东西时，我给自己定的目标很明确——不追求复现一个 GPT，而是把一条最小可用的训练加推理链路，用最朴素的方式搭出来，每一行都能解释清楚它在干什么。下面这套东西，就是我踩了无数坑之后沉淀下来的完整思路和实操记录，代码能跑、参数能算、问题能查，你可以直接抄作业。

2. 整体架构设计与技术选型思路

2.1 为什么不用现成框架，而是选择手写核心层

很多人第一反应是：都什么年代了还手写矩阵乘法，这不是重复造轮子吗？我一开始也这么想，直到有一次线上模型输出突然全变成同一个值，我查了整整两天，最后发现是某个框架版本里softmax的数值稳定实现变了。那一刻我才意识到，如果你不理解底层在干什么，你连问题出在哪一层都定位不了。

从零实现的价值不在于替代框架，而在于建立“心智模型”。我的设计思路是这样的：用 NumPy 手写前向和反向传播，把每一层的梯度推导都显式写出来；然后用 PyTorch 做同样的计算，逐层对比数值，确认我手写的梯度是对的；最后再把性能敏感的部分（比如卷积、矩阵乘）换成优化实现。这样走一遍，框架里那些autograd、cuda、amp的魔法，就全变成你能看懂的普通代码了。

选型上我坚持几个原则。第一，数值计算用 NumPy 起步，因为它足够慢，慢到你能看清每一步的内存分配和计算顺序，这对理解广播机制和维度变换极其重要。第二，对比验证用 PyTorch，不是因为它最好，而是因为它生态最全，遇到问题搜得到答案。第三，性能优化阶段才引入底层加速，比如用numba做 JIT，或者用einsum替代手写循环。这个渐进式的路线，能让你在每一步都知道自己为什么慢、慢在哪、怎么快起来。

提示：不要一上来就追求速度。我见过太多人直接上 CUDA 手写 kernel，结果连反向传播的链式法则都没搞明白，最后调出来的梯度是错的，训练 loss 不降反升，白白浪费一周。

2.2 模块划分：从张量到训练循环的六层结构

我把整个工程拆成了六个层次，每一层只依赖它下面的一层，这样调试的时候可以逐层隔离问题。这个划分不是拍脑袋想的，而是根据“出错时最容易定位”的原则反推出来的。

这么分的好处是，当你的 loss 变成 NaN 时，你可以先冻结 L6，用固定输入测 L5 的损失函数；如果损失函数没问题，再往下测 L4 的某一层输出；一层层剥，总能找到那个产生 inf 或 nan 的算子。我实测下来，这套分层排查法能把定位时间从平均半天压缩到二十分钟以内。

2.3 性能与可读性的平衡策略

从零实现最大的矛盾就是：写得越清楚，跑得越慢；跑得越快，代码越像天书。我的策略是双版本并行。每个核心算子我都写两个版本：一个naive版本，用最直白的循环，注释写满每一步的数学含义；一个fast版本，用向量化、einsum、甚至numba加速。然后用单元测试保证两个版本输出一致。

这样做的好处是，当你怀疑 fast 版本有 bug 时，可以拿 naive 版本当“标准答案”来对。我印象最深的是实现卷积反向传播那次，fast 版本梯度总是差一点点，最后用 naive 版本逐元素对比，发现是 padding 边界处理时多算了一个像素。这种错误，光看 fast 版本的代码是根本看不出来的。

3. 核心细节解析与手写实操要点

3.1 张量类的设计：存储、步长与广播的底层逻辑

张量看起来简单，其实坑最多。我的Tensor类只存三样东西：data（一个扁平的 NumPy 数组）、shape（逻辑形状）、strides（每个维度跨多少元素）。为什么不直接存多维数组？因为扁平存储加步长，才能让你真正理解 view、transpose、broadcast 这些操作到底有没有拷贝数据。

举个例子，一个形状为(3, 4)的矩阵，按行优先存储，它的 strides 是(4, 1)。当你做transpose变成(4, 3)时，你不需要动 data，只需要把 strides 改成(1, 4)。这就是为什么转置几乎不花时间。而当你做reshape成(2, 6)时，如果内存布局允许，同样只改 strides；如果不允许，才需要真正拷贝。理解这一点，你就能明白为什么有些操作快、有些操作慢。

广播是另一个重灾区。规则本身不复杂：从最后一个维度往前对齐，维度为 1 的可以扩展。但实现的时候，我建议你把广播后的 strides 里对应被扩展的维度设成 0。这样在计算元素地址时，index * 0 = 0，自然就实现了“同一个值被重复读取”，不需要真的复制数据。这个技巧我在实际写的时候用了，内存占用直接降了一个数量级。

注意：手写张量时，一定要在每次形状变换后校验strides和shape是否自洽。我踩过的坑是，连续做两次 transpose 后忘了更新 strides，结果数据读出来全是乱的，但程序不报错，只是结果悄悄错了。这种 silent bug 最要命。

3.2 反向传播的手写实现：计算图与链式法则的落地

自动微分的核心是计算图。我的实现很土：每个 Tensor 有一个grad字段，还有一个_backward函数。前向计算时，每个算子负责把自己的_backward注册到输出 Tensor 上，并记录它依赖哪些输入。反向传播时，从 loss 开始，按拓扑逆序依次调用_backward，把梯度累加到输入的grad上。

这里最关键的是梯度累加而不是覆盖。因为一个 Tensor 可能被多个下游算子使用，比如x同时参与了y = x * 2和z = x + 3，那么x.grad应该是两条路径梯度之和。我一开始写成覆盖，结果训练完全不收敛，查了半天才发现是梯度被后一次计算冲掉了。

链式法则的落地，我用一个具体例子说明。假设y = relu(Wx + b)，反向时：

• 先算relu的梯度：输入大于 0 的地方梯度为 1，否则为 0，逐元素乘上上游梯度。
• 再算加法+b的梯度：对b的梯度就是上游梯度按 batch 维度求和。
• 再算矩阵乘Wx的梯度：对W的梯度是上游梯度^T · x，对x的梯度是上游梯度 · W^T。

每一步我都用数值梯度验证过：取一个很小的eps，用(f(x+eps) - f(x-eps)) / (2*eps)算近似梯度，和反向传播的结果对比，误差在1e-6以内才算通过。这个验证步骤绝对不能省，它是你信心的唯一来源。

3.3 数值稳定性处理：softmax、log 与梯度裁剪

数值稳定性是从零实现里最容易翻车的地方。最典型的就是softmax。直接按定义算exp(x) / sum(exp(x))，当x里有比较大的值时，exp会溢出成 inf，然后 inf 除以 inf 就是 nan。解决办法是减去最大值：先算x - max(x)，这样指数最大是 0，不会溢出，而且数学上结果完全等价。

交叉熵损失也有同样的问题。如果你先算 softmax 再取 log，中间那步log(0)会变成负无穷。正确做法是把 softmax 和 log 合并成一个算子，用 log-sum-exp 技巧直接算，避免中间结果溢出。我在实现的时候，专门写了一个log_softmax，然后交叉熵直接基于它算，这样既快又稳。

梯度裁剪也是必备的。尤其是 RNN 或者深层网络，梯度爆炸会让参数瞬间变成 nan。我的做法是全局范数裁剪：算所有参数梯度的 L2 范数，如果超过阈值（我一般设 1.0 或 5.0），就整体缩放。这个操作放在优化器step之前，几乎不增加计算量，但能救命。

提示：判断是否发生梯度爆炸，不用等 loss 变 nan。你可以在训练循环里打印梯度范数，如果它突然从个位数跳到几百上千，那就是前兆，赶紧加裁剪或者调小学习率。

4. 完整实操流程与关键环节实现

4.1 环境搭建与依赖管理

环境这块我不建议用太花哨的工具，venv加pip就够了。核心依赖只有三个：numpy做数值计算，pytest做单元测试，matplotlib画训练曲线。如果你要对比 PyTorch，再加一个torch，但注意别在核心代码里 import 它，只在测试文件里用。

我的目录结构是这样的：

ai-engineering-from-scratch/ ├── core/ │ ├── tensor.py # L1 张量基础 │ ├── ops.py # L2 算子 │ ├── autograd.py # L3 自动微分 │ └── nn.py # L4 网络层 ├── optim/ │ ├── sgd.py │ └── adam.py ├── data/ │ └── loader.py ├── tests/ │ ├── test_tensor.py │ ├── test_grad.py │ └── test_compare_torch.py └── train.py

这个结构的好处是，core里完全不依赖任何训练逻辑，你可以单独 import 出来做实验。tests里每个文件对应一层，跑pytest就能快速定位是哪一层出了问题。

4.2 手写一个线性层并验证梯度

我从最简单的线性层开始。前向是y = x @ W + b，其中x形状(batch, in_features)，W形状(in_features, out_features)，b形状(out_features,)。反向时，对W的梯度是x.T @ grad_y，对b的梯度是grad_y.sum(axis=0)，对x的梯度是grad_y @ W.T。

参数初始化我用的是 Xavier：W从均值为 0、标准差为sqrt(2 / (in_features + out_features))的正态分布采样。为什么用这个？因为如果初始化太大，前向输出会爆炸；太小，输出会趋近于 0，梯度也跟着消失。Xavier 让每一层的输出方差大致保持不变，训练初期更稳。

验证梯度的时候，我写了一个通用函数：

def numerical_grad(f, x, eps=1e-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 fx1 = f(x) x[idx] = old - eps fx2 = f(x) x[idx] = old grad[idx] = (fx1 - fx2) / (2 * eps) it.iternext() return grad

这个函数虽然慢，但它是你验证反向传播正确性的黄金标准。我每写一个新算子，都会用它跑一遍，误差超过1e-5就说明有问题。

4.3 训练循环与数据管道的搭建

训练循环看起来简单，但细节决定成败。我的循环长这样：

for epoch in range(num_epochs): for x_batch, y_batch in dataloader: y_pred = model(x_batch) loss = cross_entropy(y_pred, y_batch) model.zero_grad() loss.backward() clip_grad_norm(model.params, max_norm=1.0) optimizer.step() metrics.update(loss.item(), y_pred, y_batch) print(f"epoch {epoch}, loss {metrics.avg_loss:.4f}, acc {metrics.acc:.4f}")

这里有几个关键点。第一，zero_grad必须在backward之前调用，否则梯度会跨 batch 累加，相当于变相增大了 batch size，学习率就不匹配了。第二，clip_grad_norm放在backward之后、step之前，顺序不能乱。第三，数据管道我用了一个简单的生成器，每次 yield 一个 batch，支持 shuffle。shuffle 的实现是每个 epoch 开始时生成一个随机排列的索引数组，然后按这个顺序取数据。

数据管道的性能也很重要。如果你的数据加载比模型计算还慢，GPU 就会一直空转。我的做法是提前把数据转成 NumPy 数组，加载时只做切片和转置，避免在循环里做复杂的预处理。如果数据量特别大，可以用多进程预取，但那是后话，先把单进程跑通再说。

4.4 从 NumPy 到 PyTorch 的逐层对比验证

这一步是整个流程里最有价值的。我会用同样的输入、同样的参数初始化，分别跑我的实现和 PyTorch，然后逐层对比输出。对比的顺序是：先比前向输出，再比损失值，最后比每一层的梯度。

前向对比时，误差应该在1e-6量级。如果某一层误差突然变大，说明那一层的实现有问题。梯度对比时，我会把 PyTorch 的param.grad和我的param.grad拉平后算余弦相似度，正常应该接近 1.0。如果相似度只有 0.9 甚至更低，说明梯度方向有偏差，训练可能还能跑，但收敛会变慢。

我印象最深的一次，是发现我的LayerNorm反向梯度和 PyTorch 差了 5%。查了半天，发现是我在计算方差时用了ddof=0，而 PyTorch 用的是有偏估计，但反向传播里对均值和方差的梯度处理有细微差别。这种差异，只有逐层对比才能发现。

5. 常见问题与排查技巧实录

5.1 训练不收敛的五大元凶

训练不收敛是最常见的问题，我把它归结为五类原因，按排查优先级排列：

我踩过最坑的一次是“标签错位”。数据管道里做 shuffle 时，我分别 shuffle 了特征和标签，结果两者对不上，模型学了个寂寞。这种问题不会报错，loss 也会缓慢下降（因为模型在拟合随机噪声），但准确率永远上不去。后来我改成用一个索引数组同时 shuffle 特征和标签，问题就解决了。

5.2 梯度消失与爆炸的定位方法

梯度问题的定位，核心是在每一层后面打印梯度范数。我的做法是在backward之后，遍历所有参数，打印它们的梯度 L2 范数。如果发现从某一层开始，梯度范数骤降到1e-8以下，那就是梯度消失；如果骤增到1e3以上，那就是爆炸。

梯度消失通常和激活函数有关。Sigmoid 在输入绝对值大时梯度趋近于 0，深层网络里连乘下来就没了。解决办法是换 ReLU 或者用残差连接。梯度爆炸则多和初始化、学习率有关，加梯度裁剪是最直接的缓解手段。

还有一个隐蔽的问题是梯度为 0。如果你发现某一层梯度全是 0，先检查这一层是不是被detach了，或者它的输入是不是常量。我有一次写了个自定义算子，忘了把输入注册到计算图里，结果那一层永远不更新，查了好久才发现。

5.3 内存与显存溢出的排查思路

内存问题分两种：一种是真不够，一种是泄漏。真不够的话，减小 batch size 或者用梯度累积是最快的办法。梯度累积就是跑几个小 batch，把梯度累加起来，再统一step一次，效果等价于大 batch，但内存占用小。

泄漏的话，通常是计算图没释放。比如你在循环里不断把 loss 存进一个列表，而 loss 又连着整个计算图，那内存就会一直涨。解决办法是存loss.item()而不是loss本身，或者在不需要梯度的地方用with torch.no_grad()。

我实测下来，一个中等规模的模型，如果 batch size 是 64，输入是(64, 3, 224, 224)，光输入数据就占64*3*224*224*4字节，约 38MB。如果中间特征图没及时释放，很容易就上 G 了。所以我的习惯是，每写完一个模块，就用一个小 batch 跑一遍，看内存峰值是多少，心里有个数。

5.4 数值精度问题的隐蔽陷阱

浮点数不是精确的，这一点在从零实现里体现得淋漓尽致。最典型的是累加顺序不同导致结果不同。比如sum一个很大的数组，从左往右加和从右往左加，结果可能差1e-7。这个误差在单层看不出来，但深层网络里会累积。

我的应对策略是：关键计算用 float64，训练用 float32。验证梯度的时候，全部转成 float64，这样数值误差最小，容易发现真正的逻辑错误。训练的时候用 float32，省内存也够用。另外，比较两个浮点数时，永远不要用==，要用abs(a - b) < tol，tol根据量级取1e-5到1e-7。

还有一个坑是log(0)。交叉熵里如果预测概率是 0，log 就是负无穷，loss 直接变 inf。解决办法是在 log 之前加一个极小值eps，比如1e-12，或者直接用log_softmax的稳定实现。我两种都用过，后者更干净。

注意：如果你在 GPU 上跑，float32 的精度问题会比 CPU 更明显，因为 GPU 的某些算子为了速度会牺牲精度。遇到诡异的不一致，先怀疑精度，再怀疑逻辑。

6. 性能优化与工程化收尾

6.1 向量化替代循环的实操技巧

从零实现最容易写成三重循环，跑起来慢得让人想砸键盘。向量化的核心思想是：把循环维度变成数组维度，让 NumPy 在 C 层面帮你循环。比如计算y = x @ W + b，如果你用循环逐元素算，复杂度是O(batch * in * out)，但用np.dot或者@，底层是 BLAS 库，速度能快几十倍。

我常用的向量化技巧有三个。第一，用einsum表达复杂的多维乘法，它可读性好，而且 NumPy 会帮你优化。第二，用广播替代显式复制，比如x[:, None, :] * y[None, :, :]，避免tile带来的内存开销。第三，用np.add.at处理不规则的累加，比如 embedding 的反向传播，比循环快很多。

实测下来，一个原本需要 30 秒的 epoch，向量化之后能压到 2 秒以内。这个提升不是线性的，因为向量化还减少了 Python 解释器的开销。

6.2 用 numba 做 JIT 加速的边界

numba能把 Python 函数编译成机器码，对数值循环特别有效。但它有边界：不支持所有 NumPy 特性，对动态类型不友好，编译本身也有开销。我的经验是，只对热点循环用 numba，比如卷积的滑动窗口计算，或者自定义的损失函数。

用的时候注意两点。第一，函数里尽量只用基本类型和 NumPy 数组，别用列表、字典这些。第二，第一次调用会触发编译，慢是正常的，从第二次开始才快。我一般会先跑一个 warmup，把编译时间排除在计时之外。

如果 numba 搞不定，退而求其次用cython或者直接写 C 扩展，但那是最后的手段。大多数情况下，好的向量化加上 numba，已经能覆盖 90% 的性能需求。

6.3 单元测试与持续验证的落地

从零实现最怕的就是“改一处、崩三处”。我的做法是给每个算子写单元测试，测试内容包括：形状正确性、数值正确性（对比数值梯度）、边界情况（空输入、单元素、极端值）。然后用pytest一键跑全部测试，每次改完代码先跑测试，通过了再训练。

测试文件里我会用pytest.mark.parametrize覆盖多组参数，比如不同的 batch size、不同的特征维度。这样能发现很多隐藏的维度 bug。我还会写一个test_compare_torch.py，专门对比我的实现和 PyTorch 的输出，确保数值一致。

这套测试体系建立起来之后，我改代码的胆子大多了。因为我知道，只要测试全绿，核心逻辑就是对的，剩下的只是调参和优化。

6.4 从实验代码到可复用模块的整理

最后一步是把实验代码整理成可复用的模块。我的原则是：接口清晰、依赖最少、文档到位。每个模块只暴露必要的类和函数，内部实现细节用下划线前缀标记为私有。依赖关系用requirements.txt固定版本，避免环境漂移。

文档我写得不多，但每个公开函数都有一行 docstring，说明输入输出形状和数学含义。比如linear(x, W, b)的 docstring 会写：x: (batch, in), W: (in, out), b: (out,), return: (batch, out)。这样别人用的时候不用猜维度。

整理完之后，我把整个项目打成一个包，可以在其他项目里pip install -e .直接引用。这样下次做新实验，就不用重新搭一遍了。

这个项目我断断续续做了大概两个月，中间推翻重来了三次。最大的收获不是代码本身，而是那种“我知道每一行在干什么”的踏实感。后来线上再遇到模型异常，我第一反应不再是“换个版本试试”，而是“我去看看那一层的梯度”。这种从被动到主动的转变，才是我觉得从零实现最大的价值。如果你也在做类似的事情，我的建议是：别急，一层一层来，每写一个算子就用数值梯度验一遍，慢就是快。