从零构建AI工程能力：手写自动微分与Transformer推理引擎实战

1. 从零搭建AI工程能力：为什么我劝你别再当“调包侠”

这两年AI岗位的招聘需求翻了何止三倍，但真正能扛住面试官追问的人少得可怜。我面过不下五十个候选人，简历上清一色写着“熟悉PyTorch、TensorFlow”，结果一问“手写一个带反向传播的两层神经网络”，能当场在白板上推明白的不到两成。这就是“ai-engineering-from-scratch”这个项目标题背后最扎心的现实——绝大多数人所谓的AI工程能力，其实只是熟练使用几个高层API，一旦底层出问题，连日志都看不懂。

所谓“from scratch”，不是让你从晶体管开始造计算机，而是从数学原理、数据结构、训练循环、工程部署这四个层面，把AI系统的每一块积木都亲手搭一遍。这个项目适合三类人：第一类是刚入行的算法工程师，想摆脱“只会调参”的标签；第二类是有后端或数据背景的开发者，想转AI方向但被各种框架的抽象层绕晕；第三类是在校学生，课程只教了理论，没告诉你工业界真正跑一个模型要踩多少坑。

我花了大概三个月时间，把这条路径完整走了一遍，从用NumPy实现第一个线性回归，到用纯Python写一个能跑通的Transformer推理引擎，再到把模型封装成可观测的服务。整个过程踩的坑比我想象的多得多，但也正是这些坑，让我真正理解了什么叫“工程能力”。下面我把这套方法拆开揉碎讲清楚，你照着做，至少能少走半年弯路。

2. 整体设计思路：为什么我选择“自底向上”而不是“自顶向下”

2.1 主流学习路径的致命缺陷

市面上大部分AI教程走的是“自顶向下”路线：先教你调sklearn的fit和predict，再教你用Keras搭几层网络，最后告诉你“底层原理很复杂，先会用就行”。这种路径的问题在于，你永远在黑盒外面打转。模型不收敛，你不知道是学习率问题还是梯度消失；推理延迟高，你不知道是算子实现问题还是内存布局问题。

我试过带一个新人用这种方式入门，三个月后他能跑通图像分类的demo，但当我让他把模型部署到边缘设备上、把延迟从200ms压到50ms时，他完全无从下手。因为他不知道卷积在底层是怎么算的，不知道内存复用和算子融合是什么，更不知道量化到底损失了什么。

“from scratch”的思路完全反过来：先理解最小可运行单元，再逐层向上抽象。就像学开车，你得先知道离合器、油门、变速箱各自干什么，再去学怎么在高速上超车。这个顺序不能反。

2.2 四层能力模型的设计逻辑

我把整个学习路径拆成四个层次，每一层都有明确的产出物和验收标准：

这个分层的关键在于每一层都必须能独立运行。你写自动微分引擎时不需要任何深度学习框架，写推理引擎时不需要PyTorch的nn.Module。这种“断奶式”训练，才能真正逼你理解每个环节的本质。

2.3 为什么不用现成框架“抄近路”

有人会问：既然PyTorch已经这么成熟了，为什么还要自己造轮子？我的回答是：造轮子不是为了用，而是为了懂。你不需要在生产环境用自己写的自动微分，但你必须知道PyTorch的autograd在背后做了什么，否则你连loss.backward()为什么报“梯度为None”都搞不清楚。

我自己的经验是，手写一遍反向传播之后，再看PyTorch的retain_graph参数、detach()方法、torch.no_grad()上下文，瞬间就通透了。这些概念在文档里看十遍，不如自己踩一次“计算图被释放”的坑来得深刻。

3. 核心细节解析：从零实现自动微分引擎的完整思路

3.1 计算图到底该怎么建

自动微分的核心是计算图。每个张量是图上的一个节点，每个运算是一条边。前向传播时记录运算关系，反向传播时沿着图反向求导。听起来简单，但实现起来有几个关键决策。

第一个决策是用静态图还是动态图。静态图先定义再运行，适合优化但调试困难；动态图边运行边建图，调试友好但性能略差。我选择动态图，因为学习阶段调试便利性远比性能重要。PyTorch早期能打败TensorFlow，靠的就是动态图的易用性。

第二个决策是节点存储什么。每个节点需要存：数据（data）、梯度（grad）、产生该节点的运算（_op）、父节点列表（_parents）、以及反向传播函数（_backward）。这里有个容易忽略的点：梯度必须累加而不是覆盖。因为一个张量可能被多个下游节点使用，反向传播时它的梯度是各分支梯度之和。我一开始写成覆盖，结果多分支网络训练完全不收敛，排查了一整天才发现这个问题。

class Tensor: def __init__(self, data, _parents=(), _op=''): self.data = np.array(data, dtype=np.float64) self.grad = np.zeros_like(self.data) self._backward = lambda: None self._parents = _parents 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

上面这段代码里，self.grad += out.grad就是累加的关键。如果你写成self.grad = out.grad，多分支场景下梯度就错了。

3.2 反向传播的拓扑排序为什么必不可少

反向传播必须从输出往输入方向、按拓扑序依次调用每个节点的_backward函数。如果你随便遍历节点，可能出现某个节点的梯度还没从所有下游收集完，就提前传播给上游了，结果就是梯度不完整。

拓扑排序的实现有两种：一种是DFS后序遍历，一种是Kahn算法。我推荐DFS后序，代码更简洁：

def backward(self): topo = [] visited = set() def build_topo(v): if v not in visited: visited.add(v) for parent in v._parents: build_topo(parent) topo.append(v) build_topo(self) self.grad = np.ones_like(self.data) for node in reversed(topo): node._backward()

这里有个细节：输出节点的初始梯度设为1。因为∂L/∂L = 1，这是链式法则的起点。很多人会忘记这一步，导致所有梯度都是0。

3.3 广播机制带来的梯度求和陷阱

NumPy的广播机制让不同形状的张量能直接运算，但反向传播时必须把广播后的梯度还原回原始形状。比如形状(3,4)的张量和形状(4,)的偏置相加，前向传播时偏置被广播成(3,4)，反向传播时偏置的梯度应该是(3,4)梯度在第一个维度上求和，得到(4,)。

这个坑我踩过两次。第一次是偏置梯度形状不对，第二次是批量矩阵乘法时梯度维度对不上。解决方法是在每个运算的_backward里显式处理形状还原：

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 += _unbroadcast(out.grad * other.data, self.data.shape) other.grad += _unbroadcast(out.grad * self.data, other.data.shape) out._backward = _backward return out def _unbroadcast(grad, shape): while grad.ndim > len(shape): grad = grad.sum(axis=0) for i, dim in enumerate(shape): if dim == 1: grad = grad.sum(axis=i, keepdims=True) return grad

_unbroadcast这个函数是通用解法，任何涉及广播的运算都可以复用它。写一次，后面所有运算都省心。

注意：梯度累加前必须确保形状一致，否则NumPy会静默广播，导致梯度悄悄出错。建议在开发阶段加断言检查形状。

4. 实操过程：手写两层神经网络完成MNIST分类

4.1 数据加载与预处理的关键细节

MNIST数据集虽然简单，但预处理有几个容易翻车的点。第一是归一化：像素值从0-255缩放到0-1，否则梯度爆炸。第二是标签格式：如果用交叉熵损失，标签应该是one-hot编码；如果用稀疏交叉熵，标签是整数索引。我选择one-hot，因为手写交叉熵更直观。

第三是数据打乱：每个epoch开始前必须打乱训练集顺序，否则模型会学到样本顺序的虚假相关性。我一开始忘了打乱，训练准确率卡在60%上不去，排查了半天才发现是数据顺序问题。

def load_mnist(): from sklearn.datasets import fetch_openml X, y = fetch_openml('mnist_784', version=1, return_X_y=True, as_frame=False) X = X.astype(np.float64) / 255.0 y = y.astype(np.int64) X_train, X_test = X[:60000], X[60000:] y_train, y_test = y[:60000], y[60000:] y_train_onehot = np.zeros((len(y_train), 10)) y_train_onehot[np.arange(len(y_train)), y_train] = 1 return X_train, y_train_onehot, X_test, y_test

4.2 网络结构设计与参数初始化

两层网络的结构是：输入784维 → 隐藏层128维（ReLU激活）→ 输出10维（Softmax）。参数量计算：784×128 + 128 + 128×10 + 10 = 101770个参数。这个规模用纯NumPy训练完全可行，一个epoch大概30秒。

参数初始化用He初始化，因为ReLU激活函数会把负半轴梯度清零，如果用标准正态分布初始化，深层网络的方差会逐层衰减。He初始化把权重标准差设为√(2/fan_in)，保证前向传播时每层输出方差稳定。

def init_params(): W1 = np.random.randn(784, 128) * np.sqrt(2.0 / 784) b1 = np.zeros(128) W2 = np.random.randn(128, 10) * np.sqrt(2.0 / 128) b2 = np.zeros(10) return W1, b1, W2, b2

这里b1和b2初始化为0是安全的，因为ReLU和Softmax对偏置的梯度不会消失。但如果你用Sigmoid激活，偏置也建议小随机初始化。

4.3 前向传播与损失计算的数值稳定性

Softmax + 交叉熵的组合有个经典陷阱：指数运算溢出。如果logits里有大于88的值，np.exp就会返回inf。解决方法是在Softmax之前减去每行的最大值：

def softmax(x): x_shifted = x - np.max(x, axis=1, keepdims=True) exp_x = np.exp(x_shifted) return exp_x / np.sum(exp_x, axis=1, keepdims=True)

这个操作数学上等价，因为Softmax对输入平移不变，但数值上稳定得多。交叉熵损失也要用log_softmax的形式计算，避免先算Softmax再取log导致的精度损失：

def cross_entropy_loss(logits, y_onehot): log_probs = logits - np.max(logits, axis=1, keepdims=True) log_probs = log_probs - np.log(np.sum(np.exp(log_probs), axis=1, keepdims=True)) return -np.mean(np.sum(y_onehot * log_probs, axis=1))

4.4 反向传播的逐层推导与实现

反向传播的推导必须自己动手推一遍，否则永远记不住。以两层网络为例：

输出层梯度：∂L/∂z2 = softmax(z2) - y，这是Softmax+交叉熵的经典结论，推导过程用链式法则展开即可。

隐藏层梯度：∂L/∂z1 = (∂L/∂z2 · W2ᵀ) ⊙ ReLU'(z1)，其中ReLU'在z1>0时为1，否则为0。

参数梯度：∂L/∂W2 = a1ᵀ · ∂L/∂z2，∂L/∂b2 = sum(∂L/∂z2, axis=0)。

def backward(X, y_onehot, cache): a1, z1, a2, z2 = cache m = X.shape[0] dz2 = (softmax(z2) - y_onehot) / m dW2 = a1.T @ dz2 db2 = np.sum(dz2, axis=0) da1 = dz2 @ W2.T dz1 = da1 * (z1 > 0) dW1 = X.T @ dz1 db1 = np.sum(dz1, axis=0) return dW1, db1, dW2, db2

注意dz2除以了m，因为损失是batch平均。如果你忘了这个除法，梯度会大m倍，学习率需要相应调整。

4.5 训练循环与超参数调优实录

训练循环的核心是：前向传播 → 计算损失 → 反向传播 → 更新参数。超参数方面，我实测下来比较稳的配置是：batch size 64，学习率0.1（配合SGD），训练20个epoch。学习率太大会震荡，太小收敛慢。你可以用学习率预热：前5个epoch用0.01，后面逐步升到0.1。

W1, b1, W2, b2 = init_params() lr = 0.1 batch_size = 64 for epoch in range(20): indices = np.random.permutation(len(X_train)) for i in range(0, len(X_train), batch_size): batch_idx = indices[i:i+batch_size] X_batch = X_train[batch_idx] y_batch = y_train_onehot[batch_idx] cache = forward(X_batch, W1, b1, W2, b2) loss = cross_entropy_loss(cache[3], y_batch) grads = backward(X_batch, y_batch, cache) W1 -= lr * grads[0] b1 -= lr * grads[1] W2 -= lr * grads[2] b2 -= lr * grads[3]

我实测这个配置在MNIST测试集上能到97.5%的准确率。如果你想更高，可以加动量或者用Adam，但那就偏离“from scratch”的初衷了。

实操心得：训练初期loss不下降，先检查数据预处理和梯度形状；训练中期loss震荡，调小学习率；训练后期loss不降，可能是过拟合，加Dropout或L2正则。

5. 从训练到推理：手写Transformer推理引擎的核心要点

5.1 为什么推理和训练是两码事

训练时你关心的是梯度能不能算对、loss能不能降下去；推理时你关心的是延迟、吞吐、内存。这两个阶段的优化目标完全不同。训练可以用大batch、混合精度、梯度累积；推理必须考虑KV Cache、算子融合、量化。

我见过太多人训练模型一把梭，部署时发现推理速度慢十倍。原因就是训练和推理的计算图不一样。比如Dropout只在训练时启用，BatchNorm在推理时用滑动平均，Attention在推理时可以用KV Cache避免重复计算。

5.2 KV Cache的原理与内存计算

自回归生成时，每生成一个token都要重新计算所有历史token的Key和Value，这是巨大的浪费。KV Cache的思路是：把历史token的K和V缓存起来，新token只需要计算自己的Q、K、V，然后和缓存的K、V做Attention。

内存占用计算：假设模型有L层，每层有H个注意力头，每个头维度为D，序列长度为S，batch size为B，精度为FP16（2字节）。KV Cache总内存 = 2 × L × H × D × S × B × 2字节。以GPT-2为例，L=12，H=12，D=64，S=1024，B=1，算下来约37MB。序列越长内存越大，这就是为什么长上下文推理这么吃显存。

class KVCache: def __init__(self, num_layers, num_heads, head_dim, max_seq_len): self.cache = [ { 'k': np.zeros((num_heads, max_seq_len, head_dim)), 'v': np.zeros((num_heads, max_seq_len, head_dim)) } for _ in range(num_layers) ] self.seq_len = 0 def update(self, layer_idx, new_k, new_v): self.cache[layer_idx]['k'][:, self.seq_len] = new_k self.cache[layer_idx]['v'][:, self.seq_len] = new_v

5.3 Attention的纯NumPy实现

Attention的核心公式是softmax(QKᵀ/√d)V。实现时有几个细节：第一是缩放因子√d不能忘，否则点积结果太大会让Softmax饱和；第二是mask处理，生成时不能看到未来的token；第三是数值稳定性，Softmax前要减最大值。

def attention(Q, K, V, mask=None): d_k = Q.shape[-1] scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k) if mask is not None: scores = scores + mask * -1e9 scores = scores - np.max(scores, axis=-1, keepdims=True) weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True) return weights @ V

mask通常是一个上三角矩阵，未来位置为负无穷，Softmax后权重为0。注意-1e9而不是-inf，因为-inf参与运算会产生NaN。

5.4 文本生成策略与采样技巧

生成文本时，最简单的策略是贪心搜索：每次选概率最大的token。但贪心搜索容易生成重复、无聊的文本。更好的策略是温度采样：把logits除以温度T，T<1让分布更尖锐，T>1让分布更平坦。T=0.7通常是比较好的平衡点。

还有Top-K采样：只从概率最高的K个token中采样，避免长尾噪声。Top-P采样（核采样）：从累积概率达到P的最小token集合中采样，动态调整候选集大小。我实测Top-P=0.9配合温度0.8，生成质量比较稳定。

def sample(logits, temperature=0.8, top_p=0.9): logits = logits / temperature probs = softmax(logits) sorted_idx = np.argsort(probs)[::-1] sorted_probs = probs[sorted_idx] cumsum = np.cumsum(sorted_probs) cutoff = np.searchsorted(cumsum, top_p) sorted_probs[cutoff+1:] = 0 sorted_probs /= np.sum(sorted_probs) chosen = np.random.choice(sorted_idx, p=sorted_probs) return chosen

注意：采样前必须把logits转成概率分布，且要处理数值稳定性。温度不能设得太低，否则退化成贪心搜索；也不能太高，否则生成乱码。

6. 工程化落地：把模型封装成可观测的推理服务

6.1 服务框架选型与性能对比

推理服务的框架选择很多：Flask简单但性能差，FastAPI异步支持好，Triton Inference Server功能全但学习曲线陡。我推荐从FastAPI起步，因为它异步支持好、自动生成API文档、生态成熟。实测下来，FastAPI在单机上的QPS能到500以上，足够中小规模场景。

如果你追求极致性能，可以考虑用ONNX Runtime或TensorRT做推理后端，FastAPI只做请求转发。但那是另一个话题了，from scratch阶段先把服务跑通再说。

6.2 请求批处理与动态合并

推理服务的吞吐瓶颈往往不在计算，而在请求调度。单个请求进来就推理一次，GPU利用率极低。解决方案是动态批处理：把短时间内到达的多个请求合并成一个batch，一起推理，再拆分结果返回。

实现上可以用一个队列加一个后台worker：请求进来先入队，worker每隔几毫秒取一批，凑够batch size或超时就推理。这样能把吞吐提升5-10倍。

import asyncio from collections import deque class BatchScheduler: def __init__(self, max_batch_size=32, max_wait_ms=10): self.queue = deque() self.max_batch_size = max_batch_size self.max_wait_ms = max_wait_ms async def add_request(self, input_data): future = asyncio.Future() self.queue.append((input_data, future)) return await future async def process_loop(self): while True: if not self.queue: await asyncio.sleep(0.001) continue batch = [] start = asyncio.get_event_loop().time() while self.queue and len(batch) < self.max_batch_size: batch.append(self.queue.popleft()) if (asyncio.get_event_loop().time() - start) * 1000 > self.max_wait_ms: break inputs = [item[0] for item in batch] results = self.model_inference(inputs) for (_, future), result in zip(batch, results): future.set_result(result)

6.3 监控指标暴露与告警配置

服务上线后必须能观测。核心指标有四个：QPS、P99延迟、错误率、GPU利用率。QPS和延迟用Prometheus的Histogram记录，错误率用Counter，GPU利用率用pynvml采集。

from prometheus_client import Histogram, Counter, start_http_server REQUEST_LATENCY = Histogram('inference_latency_seconds', 'Inference latency') REQUEST_COUNT = Counter('inference_requests_total', 'Total requests', ['status']) @app.post("/predict") async def predict(request: Request): start = time.time() try: result = await model.predict(request) REQUEST_COUNT.labels(status='success').inc() return result except Exception as e: REQUEST_COUNT.labels(status='error').inc() raise finally: REQUEST_LATENCY.observe(time.time() - start)

告警规则建议设三条：P99延迟超过200ms告警、错误率超过1%告警、GPU利用率持续低于20%告警（说明资源浪费）。

6.4 常见线上问题与排查手册

线上问题排查讲究快准狠。我整理了一份速查表：

实操心得：线上服务一定要加输入长度限制，否则一个超长请求就能把内存打满。我吃过这个亏，一个用户传了10万字的输入，直接把服务搞崩了。

7. 踩坑记录与能力进阶路线

7.1 那些让我熬夜的典型Bug

第一个坑是梯度爆炸。手写网络训练时loss突然变成NaN，排查发现是学习率太大加上没有梯度裁剪。解决方法是在反向传播后加np.clip(grad, -1, 1)，或者用更小的学习率。

第二个坑是内存泄漏。推理服务跑几个小时后内存爆满，排查发现是KV Cache没有按请求释放。每个请求的cache用完必须显式清理，否则会一直累积。

第三个坑是数值精度。用FP16推理时，某些层的输出会溢出。解决方法是关键层用FP32，或者用动态缩放（loss scaling）技术。

第四个坑是并发竞争。多个请求同时访问同一个模型实例，导致状态混乱。解决方法是每个请求用独立的session，或者加锁串行化。

7.2 从“能跑”到“跑得好”的进阶方向

当你把上面这些全部跑通后，可以往这几个方向深入：算子融合（把多个小算子合并成一个大算子，减少内存访问）、量化（INT8量化能把模型压缩4倍，速度提升2-3倍）、蒸馏（用大模型教小模型，保持效果的同时减小体积）、分布式推理（张量并行、流水线并行，支撑超大模型）。

每一个方向都够写一本书，但核心思路是一样的：理解底层原理，才能做出正确的工程取舍。这也是“from scratch”最大的价值——它给你的不是某个具体技能，而是一种看透本质的能力。

我个人在实际操作中的体会是，手写一遍之后，再看任何AI框架的源码都不再发怵。你知道每个模块在干什么，知道哪里可能出问题，知道怎么优化。这种底气，是调包永远给不了的。最后再分享一个小技巧：把你手写的每个模块都加上单元测试，尤其是梯度检查，用数值梯度验证解析梯度，能帮你省下大量调试时间。