从零手搓AI工程:NumPy实现自动微分与Transformer推理服务

发布时间:2026/11/3 4:42:38
从零手搓AI工程:NumPy实现自动微分与Transformer推理服务 1. 从零手搓AI工程为什么我不建议你直接调包第一次看到ai-engineering-from-scratch这个项目名的时候我正坐在工位上啃一个已经调了三天的模型部署问题。当时第一反应是又来了又一个“从零开始”的教程仓库。但点进去翻了翻结构我改主意了——这东西值得认真聊一聊。先说清楚它是什么。ai-engineering-from-scratch本质上是一套从底层原理出发、逐步构建完整AI工程链路的学习与实践框架。它不满足于教你model.fit()和model.predict()而是把注意力、反向传播、梯度下降、向量检索、推理服务、批处理调度这些环节一个个拆开让你用最朴素的方式把它们重新实现一遍。能做什么能让你在面试里被问到“Attention 的 QKV 到底怎么算”的时候不是背公式而是直接在白板上推能让你在线上推理延迟飙高的时候知道该去看 KV Cache 还是去看算子融合。适合谁适合那些已经会用 PyTorch 或 TensorFlow 跑通 demo但总觉得中间隔着一层雾的工程师也适合想从传统后端转 AI 工程、需要补底层认知的开发者。我见过太多人包括几年前的我自己学 AI 的路径是这样的装环境、跑通官方示例、调参、上线、出问题、抓瞎。问题出在哪出在工程链路是断的。你知道nn.Linear能用但不知道它背后是一次矩阵乘加偏置你知道向量数据库能检索但不知道 HNSW 的图是怎么一层层跳的。ai-engineering-from-scratch要解决的就是这个断层——它逼你把每一层都亲手搭一遍哪怕搭得丑、搭得慢但搭完之后整个系统的行为对你来说不再是黑盒。这篇文章我会按我自己的实践路径来拆先讲整体设计思路和选型考量再逐层拆核心细节和实操要点然后是完整跑通一个最小可用链路的记录最后把我踩过的坑和排查经验整理成速查表。全程不藏私能抄的代码和参数我都放出来。2. 整体设计与思路拆解为什么是“手搓”而不是“调包”2.1 手搓路线的核心逻辑把黑盒变成白盒市面上不缺 AI 教程缺的是让你亲手制造黑盒再打开它的教程。ai-engineering-from-scratch的设计哲学很明确先用 NumPy 实现一个只有几十行的微型自动微分引擎再用它搭一个单层 Transformer最后把这个 Transformer 包装成一个能接收 HTTP 请求的推理服务。整条链路没有任何“魔法”每一行代码你都能追溯到数学定义。为什么非要这样因为调包调久了人会形成一种危险的直觉——以为模型就是一个forward()函数输入张量输出张量中间发生了什么与我无关。这种直觉在 90% 的场景下没问题但剩下 10% 的场景恰恰是区分普通工程师和资深工程师的地方。比如线上服务突然出现显存泄漏你如果不知道计算图的生命周期和中间激活值的持有关系就只能重启服务碰运气比如推理吞吐上不去你如果不知道 batch 维度和序列维度在算子层面的交互方式就不知道该调 batch size 还是该做算子融合。我自己的体会是手搓一遍之后你看 PyTorch 源码的感觉完全不一样了。以前看nn.MultiheadAttention觉得是一团乱麻现在能一眼看出哪部分是 QKV 投影、哪部分是缩放点积、哪部分是输出投影甚至能猜到它为什么在某些维度上做了 reshape 而不是 transpose。2.2 技术选型背后的取舍为什么用 NumPy 而不是直接上 PyTorch项目里最核心的一个选型决策是底层计算用 NumPy不用 PyTorch 的 autograd。这个选择乍看很反直觉——既然最终要用 PyTorch为什么不一开始就用原因有三层。第一层是认知负担。PyTorch 的 autograd 太方便了方便到你写loss.backward()的时候根本不需要想链式法则怎么展开。但ai-engineering-from-scratch的目标是让你想清楚每一个梯度是怎么流回去的。用 NumPy 手写反向传播你会被迫处理每个算子的局部梯度、被迫考虑广播机制对梯度求和的影响、被迫理解为什么 softmax 的雅可比矩阵长那样。这些痛苦是必要的就像学开车不能直接从自动挡开始得先知道离合和档位的配合逻辑。第二层是调试可见性。NumPy 的数组就是内存里的一块连续数据你可以随时print它的形状、数值、步长。PyTorch 的张量虽然也能看但一旦涉及计算图和设备转移调试信息就变得嘈杂。手搓阶段用 NumPy出问题的时候你能精确知道是哪一步的矩阵维度对不上而不是对着一个RuntimeError: shape mismatch猜半天。第三层是性能预期的建立。当你用 NumPy 实现完一个注意力层发现处理一个 512 长度的序列要 200 毫秒你才会真正理解为什么工业界要搞 FlashAttention、为什么要做算子融合、为什么 GPU 上的内存带宽是瓶颈而不是算力。这种对性能的“痛感”是直接调 PyTorch 永远给不了的。当然项目并没有让你一直手搓下去。它的设计是分阶段的阶段一手搓 NumPy 实现阶段二迁移到 PyTorch 并对比性能差异阶段三做工程化封装。这个渐进路径很合理既建立了底层认知又不至于让你在 NumPy 里造一辈子轮子。2.3 工程链路的完整性设计从张量到服务很多“从零实现”的教程只做到模型前向传播就结束了但ai-engineering-from-scratch的野心更大——它要覆盖从张量到线上服务的完整链路。这意味着除了模型本身你还得处理数据管道怎么把原始文本变成 token id怎么做 padding 和 attention mask怎么用 Dataset 和 DataLoader 做批处理。训练循环怎么组织 epoch、batch、optimizer step怎么处理梯度累积和梯度裁剪。模型持久化怎么保存和加载权重怎么处理版本兼容。推理服务怎么把模型包装成 API怎么处理并发请求怎么做动态批处理。性能监控怎么测量延迟和吞吐怎么定位瓶颈。这个完整性是它区别于普通教程的关键。我见过太多人模型训得很好但一上线就抓瞎就是因为训练和服务之间的鸿沟没人教。这个项目把这条鸿沟填上了虽然填得不算完美但至少给了你一个可运行的参考实现。3. 核心细节解析与实操要点手搓过程中的关键环节3.1 自动微分引擎20行代码背后的数学整个项目最精妙的部分是那个微型自动微分引擎。它的核心思想是每个张量除了存数据还存一个grad和一个_backward函数。前向传播的时候每个操作记录下自己的反向传播逻辑反向传播的时候从 loss 开始按拓扑逆序调用这些_backward函数。我按项目思路实现的时候定义了一个Tensor类关键属性是data、grad、_children和_backward。加法操作的反向传播就是把上游梯度原样传给两个子节点乘法操作的反向传播是把上游梯度乘以另一个子节点的数据。这里有个容易踩的坑广播机制。如果两个张量形状不同NumPy 会自动广播但反向传播的时候梯度需要按广播的维度求和还原。我一开始没处理这个结果梯度形状对不上训练完全不收敛。class Tensor: def __init__(self, data, _children(), _op): self.data np.array(data, dtypenp.float32) self.grad np.zeros_like(self.data) self._backward lambda: None self._prev set(_children) self._op _op def __add__(self, other): other other if isinstance(other, Tensor) else Tensor(other) out Tensor(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, Tensor) else Tensor(other) out Tensor(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这段代码看起来简单但它包含了自动微分的全部核心思想。你把它扩展一下加上__pow__、__matmul__、relu、softmax就能搭出一个完整的神经网络。我实测下来用这个引擎训练一个两层 MLP 做 MNIST 分类准确率能到 97% 左右虽然比 PyTorch 慢十倍但每一步梯度怎么来的你都清清楚楚。注意手写反向传播的时候梯度一定要用而不是。因为一个张量可能被多个下游节点使用梯度需要累加。我一开始用结果共享权重的网络完全不收敛排查了两个小时才发现是这个原因。3.2 注意力机制的实现从公式到代码的映射注意力机制是 Transformer 的核心也是手搓过程中最容易出错的地方。公式大家都背过Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V。但把它写成代码的时候有几个细节必须注意。首先是缩放因子。sqrt(d_k)里的d_k是 key 的维度不是序列长度也不是 batch size。我见过有人写成sqrt(seq_len)的训练初期 loss 直接爆炸。为什么是d_k因为 Q 和 K 的点积结果的方差与d_k成正比除以sqrt(d_k)是为了把方差拉回 1 附近避免 softmax 进入饱和区导致梯度消失。其次是mask 的处理。padding mask 和 causal mask 是两回事。padding mask 是把填充位置的注意力权重置为负无穷causal mask 是防止模型看到未来 token。实现的时候mask 要在 softmax 之前加加的是负无穷实践中用-1e9代替不是 0。我一开始把 mask 加在 softmax 之后结果 padding 位置还是有概率值模型学出来的东西完全不对。def attention(Q, K, V, maskNone): d_k Q.shape[-1] scores Q K.transpose(-2, -1) / np.sqrt(d_k) if mask is not None: scores scores mask # mask 中需要屏蔽的位置是 -1e9 weights softmax(scores, axis-1) return weights V还有一个容易忽略的点是多头注意力的维度变换。假设d_model512num_heads8那么每个头的维度是 64。你需要把(batch, seq, 512)的输入 reshape 成(batch, seq, 8, 64)再 transpose 成(batch, 8, seq, 64)然后对每个头独立做注意力最后再拼回去。这个 reshape 和 transpose 的顺序很容易搞混我建议你拿一张纸画一下维度变化比在脑子里想靠谱得多。3.3 训练循环的工程细节梯度累积与学习率调度手搓完模型结构下一步是训练循环。这部分看起来简单但工程细节很多。ai-engineering-from-scratch里重点讲了两块梯度累积和学习率调度。梯度累积是为了在显存有限的情况下模拟大 batch。做法是前向传播和反向传播跑 N 次小 batch梯度累加第 N 次之后再optimizer.step()并清零梯度。这里有个坑如果你用了 BatchNorm梯度累积会导致统计量不准因为 BatchNorm 是在每个小 batch 上算均值和方差的。解决办法是改用 LayerNorm或者确保累积的 batch 来自同一分布。学习率调度我推荐用cosine decay with warmup。warmup 是在训练初期把学习率从 0 线性增加到峰值避免初期梯度太大导致训练不稳定cosine decay 是在后期把学习率按余弦曲线降到接近 0帮助模型收敛到更平坦的极小值。具体参数上warmup 步数一般设为总步数的 5% 到 10%峰值学习率根据 batch size 调整我实测3e-4到1e-3之间比较稳。def get_lr(step, warmup_steps, total_steps, peak_lr): if step warmup_steps: return peak_lr * step / warmup_steps progress (step - warmup_steps) / (total_steps - warmup_steps) return peak_lr * 0.5 * (1 np.cos(np.pi * progress))提示梯度裁剪是另一个必备技巧。我习惯在optimizer.step()之前做clip_grad_norm_把梯度范数限制在 1.0 以内。这个操作能有效防止梯度爆炸尤其是在训练深层网络的时候。3.4 推理服务的封装从模型到 API模型训好之后怎么把它变成一个能对外提供服务的 APIai-engineering-from-scratch给了一个基于 FastAPI 的最小实现。核心思路是启动的时候加载模型权重到内存收到请求的时候做 tokenization、前向传播、detokenization返回结果。这里有几个工程要点。第一是模型加载只做一次不要每个请求都加载一遍否则延迟会高得离谱。第二是批处理如果并发请求多可以把多个请求攒成一个 batch 一起推理吞吐能提升好几倍。第三是超时控制推理服务一定要设超时否则一个慢请求可能拖垮整个服务。from fastapi import FastAPI import numpy as np app FastAPI() model None app.on_event(startup) def load_model(): global model model load_weights(model.npz) app.post(/predict) def predict(text: str): tokens tokenize(text) logits model.forward(tokens) return {result: detokenize(logits.argmax(-1))}这个实现很粗糙但足够让你理解推理服务的骨架。生产环境还需要考虑并发控制、内存管理、监控告警等但那是下一步的事。4. 实操过程与核心环节实现跑通一个最小可用链路4.1 环境准备与依赖安装我用的环境是 Python 3.10 NumPy 1.24 FastAPI 0.100。不需要 GPUCPU 就能跑通全部流程只是训练会慢一些。依赖安装很简单pip install numpy fastapi uvicorn pydantic如果你要用 PyTorch 做对比实验再加一个pip install torch。我建议先纯 NumPy 跑一遍理解每一步在干什么再迁移到 PyTorch 看性能差异。4.2 数据准备从原始文本到训练样本我用了一个小型的文本分类数据集做实验大概 5000 条样本二分类任务。数据准备的流程是读取原始文本、构建词表、把文本转成 token id 序列、做 padding、划分训练集和验证集。词表构建我用了最简单的按频率截断统计所有 token 的出现次数保留前 10000 个其余映射到unk。padding 长度设为 128超过的截断不足的补 0。这里有个细节padding 的 id 用 0但 0 也可能是一个真实 token 的 id所以需要额外维护一个 attention mask 来区分。def build_vocab(texts, max_size10000): from collections import Counter counter Counter() for text in texts: counter.update(text.split()) vocab {word: i2 for i, (word, _) in enumerate(counter.most_common(max_size))} vocab[pad] 0 vocab[unk] 1 return vocab def encode(text, vocab, max_len128): ids [vocab.get(w, 1) for w in text.split()][:max_len] mask [1] * len(ids) [0] * (max_len - len(ids)) ids ids [0] * (max_len - len(ids)) return np.array(ids), np.array(mask)4.3 模型搭建一个单层 Transformer 分类器模型结构我选了单层 Transformer encoder 加一个平均池化加线性分类头。参数量大概 200 万CPU 上训练一个 epoch 要 3 分钟左右。结构上embedding 维度 128注意力头数 4前馈网络隐藏维度 256。搭建过程中最关键的是维度对齐。我列一下各层的输入输出维度你对照着检查层输入维度输出维度Embedding(batch, seq)(batch, seq, 128)MultiHeadAttention(batch, seq, 128)(batch, seq, 128)LayerNorm(batch, seq, 128)(batch, seq, 128)FeedForward(batch, seq, 128)(batch, seq, 128)Pooling(batch, seq, 128)(batch, 128)Classifier(batch, 128)(batch, 2)每次写完一层我都会用随机数据跑一遍前向传播确认输出维度符合预期。这个习惯帮我省了很多调试时间。4.4 训练与评估损失曲线与准确率训练循环我跑了 20 个 epochbatch size 32学习率 3e-4warmup 100 步。损失从初始的 0.69 降到 0.21验证集准确率从 50% 升到 89%。这个结果不算惊艳但考虑到模型很小、数据很少已经能说明链路是通的。训练过程中我记录了每个 epoch 的损失和准确率画成曲线看收敛趋势。前 5 个 epoch 损失下降很快后面逐渐平缓没有明显的过拟合。如果你发现验证集损失开始上升那就是过拟合了需要加 dropout 或者减小模型。4.5 推理服务启动与测试训练完成后我把权重保存成 npz 文件然后用 FastAPI 启动服务。测试的时候用 curl 发请求curl -X POST http://localhost:8000/predict -H Content-Type: application/json -d {text: this is a test}返回结果正常。延迟方面单条请求大概 15 毫秒其中大部分时间花在 tokenization 和模型前向传播上。如果要做批处理可以把多个请求攒起来一起推理吞吐能提升到每秒几百条。5. 常见问题与排查技巧实录5.1 梯度不更新或更新异常这是手搓自动微分时最常见的问题。表现是 loss 不下降或者下降后突然变成 NaN。排查思路按顺序来现象可能原因排查方法loss 完全不降梯度为 0 或未累加打印每层梯度范数loss 震荡学习率太大降低学习率 10 倍试试loss 变 NaN梯度爆炸或除零加梯度裁剪检查 log 和 sqrt部分参数不更新计算图断开检查是否用了.data或detach我遇到过一次 loss 变 NaN排查了半天发现是 softmax 的实现里除了一个可能为 0 的归一化因子。加一个1e-9的 epsilon 就好了。5.2 维度不匹配的定位技巧维度错误是手搓过程中最高频的问题。我的经验是在每个操作前后打印形状。不要只在出错的地方打印要在整个前向传播链路的关键节点都打印。这样你能一眼看出是从哪一步开始维度对不上的。另外善用assert语句。在写每一层的时候加一句assert x.shape[-1] expected_dim出错的时候能立刻定位。这个习惯看起来笨但比事后调试高效得多。5.3 推理服务性能瓶颈的排查推理服务上线后如果延迟高按这个顺序排查先看 tokenization 是不是瓶颈再看模型前向传播最后看网络传输。我实测下来小模型场景下 tokenization 经常占了一半以上的时间尤其是用 Python 做字符串处理的时候。优化方法是把 tokenization 用 C 扩展或者向量化操作重写。批处理是另一个优化点。如果 QPS 不高但延迟要求严格可以不做批处理如果 QPS 高攒批能显著提升吞吐。攒批的窗口一般设 10 到 50 毫秒太短没效果太长延迟受不了。5.4 独家避坑清单不要用np.float32之外的精度手搓阶段用 float64 会让内存翻倍速度变慢而且没必要。随机种子要固定调试的时候固定np.random.seed(42)否则每次结果不一样没法对比。保存权重用np.savez而不是picklepickle 有安全风险而且跨版本兼容性差。学习率 warmup 不要省我试过不加 warmup训练初期 loss 直接飙到 10 以上加了之后平稳下降。验证集要单独划分不要用训练集评估否则准确率虚高上线就翻车。6. 从手搓到工程化的迁移路径手搓一遍之后下一步是把它迁移到生产级框架。我的建议是分三步走第一步把 NumPy 实现翻译成 PyTorch对比两者的数值差异确认理解无误第二步用 PyTorch 的nn.Module重新组织代码加上Dataset、DataLoader、optimizer等标准组件第三步用 TorchScript 或 ONNX 导出模型做推理优化。这个迁移过程本身也是学习。你会发现 PyTorch 帮你处理了很多你手搓时忽略的细节比如设备管理、混合精度、分布式训练。但因为你手搓过你知道这些细节背后是什么不会盲目信任框架。我个人在实际操作中的体会是手搓的价值不在于你以后要手写所有代码而在于你有了判断力。当框架出问题的时候你能猜到可能是哪里当性能不达标的时候你知道该往哪个方向优化。这种判断力是调包调不出来的。

关于本文作者

来自尧图内容编辑团队

尧图内容编辑团队 内容团队

尧图内容编辑团队

本文由尧图网络内容编辑团队执笔。团队由资深项目经理、前端工程师与设计师组成,所有内容均来自亲手交付的真实项目,先讲清问题、再给出可落地的解法。尧图深耕北京网站建设十年,服务过京华建材集团、智造科技等各行业客户,把一线经验沉淀为可复用的行业观察。

  • 十年建站经验,覆盖建材、制造、服务、文创等
  • 项目经理把关选题与事实准确性
  • 工程师与设计师联合撰写专业细节
  • 统一编辑规范,保证文风与排版一致
  • 每月复盘转化数据,迭代选题方向

延伸阅读

相关资讯与近期热门内容

深度阅读推荐

建站决策前值得细读的三篇

网站改版的5个关键决策
2024-08-12

网站改版的5个关键决策

什么时候该改版、改到什么程度、如何避免流量掉光,京华建材集团改版复盘给出答案。

获取专属建站方案

看完文章,把您的行业与预算告诉我们,免费获取一份量身定制的官网建设方案与报价。

立即免费咨询