1. 从零手搓AI工程为什么我不建议你直接调包很多人一听到“AI工程”这四个字第一反应就是打开某个云平台拖几个组件、调几个API然后跑通一个Demo就觉得自己已经入门了。我刚开始接触这个方向的时候也是这么想的直到有一次线上推理服务在高峰期直接雪崩日志里全是显存溢出的报错我才意识到——只会调包的人永远不知道系统在什么边界条件下会崩。ai-engineering-from-scratch这个标题核心不在“AI”而在“from scratch”。它代表的是一种学习路径不依赖现成的高层封装从最底层的张量运算、自动求导、模型结构、训练循环、推理优化一路自己搭起来。这条路走起来慢但走完之后你对整个AI系统的掌控力是完全不同的。适合谁看适合那些已经会用PyTorch或TensorFlow跑通几个模型但一遇到性能问题、部署问题、显存问题就束手无策的工程师也适合想真正理解AI系统内部运转机制、而不是停留在“调参侠”层面的开发者。这篇文章我会按照我自己从零搭建一套AI工程体系的真实顺序来展开先讲清楚为什么必须从底层理解再拆解计算图与自动微分这个地基然后是训练循环里那些框架帮你藏起来的细节接着是推理阶段的性能优化最后聊一聊工程化落地时那些文档里不会写的坑。每一部分我都会给出可以直接上手复现的代码和参数也会解释每一个选择背后的理由。提示本文假设你有基本的Python编程能力了解NumPy的基本操作不需要你精通深度学习理论但需要你愿意动手写代码而不是只看不练。2. 计算图与自动微分整个AI工程的地基2.1 为什么框架要帮你建计算图先问一个问题当你写loss.backward()的时候PyTorch到底做了什么如果你答不上来那说明你还在“调包”阶段。计算图本质上是一个有向无环图每个节点是一个张量操作每条边表示数据依赖关系。前向传播的时候框架一边计算结果一边悄悄地把每个操作的“反向函数”记录下来形成一个反向图。反向传播的时候从loss节点出发沿着反向图链式求导把梯度一路传回去。为什么框架要费这么大劲建图因为如果没有计算图你就得手动推导每一个参数的梯度公式。一个两层全连接网络还好手推也就几十行但一个带残差连接、LayerNorm、多头注意力的Transformer手推梯度基本等于自虐。计算图的价值在于它把“求导”这件事从数学问题变成了图遍历问题框架只需要知道每个基础操作的反向函数就能自动组合出任意复杂模型的梯度。但这里有个关键点很多人忽略计算图是在前向传播过程中动态构建的。这意味着你每次调用模型图都会重新建一遍。这也是为什么PyTorch早期版本在循环神经网络上性能不如静态图框架——因为每个时间步都要重建图。后来TorchScript和编译模式的出现本质上就是在解决动态图带来的开销问题。2.2 手写一个迷你自动微分引擎光说原理没意思我们直接上手写一个能跑通的最小自动微分引擎。核心思路是定义一个Tensor类它内部保存数据、梯度、以及一个_backward函数。每个操作加、乘、矩阵乘法都会返回一个新的Tensor并且把反向传播的逻辑注册到新张量的_backward上。import numpy as np class Tensor: def __init__(self, data, requires_gradFalse): self.data np.array(data, dtypenp.float64) self.grad None self.requires_grad requires_grad self._backward lambda: None self._prev set() def __add__(self, other): other other if isinstance(other, Tensor) else Tensor(other) out Tensor(self.data other.data, self.requires_grad or other.requires_grad) out._prev {self, other} def _backward(): if self.requires_grad: self.grad out.grad if self.grad is None else self.grad out.grad if other.requires_grad: other.grad out.grad if other.grad is None else 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.requires_grad or other.requires_grad) out._prev {self, other} def _backward(): if self.requires_grad: g out.grad * other.data self.grad g if self.grad is None else self.grad g if other.requires_grad: g out.grad * self.data other.grad g if other.grad is None else other.grad g 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 np.ones_like(self.data) for v in reversed(topo): v._backward()这段代码不到五十行但它已经包含了自动微分的全部核心逻辑前向时记录反向函数反向时按拓扑序逆序调用。你可以用它跑一个简单的线性回归验证梯度是否正确。我实测下来这个迷你引擎在标量运算上的梯度精度和PyTorch完全一致误差在1e-12量级。2.3 从迷你引擎到真实框架的差距在哪里你可能会问既然五十行就能搞定自动微分那PyTorch几百万行代码在干什么差距主要在三个地方。第一是性能真实框架的算子是用CUDA写的在GPU上并行执行而我们的迷你引擎是纯Python循环慢几个数量级。第二是内存管理真实框架有内存池、有原地操作、有梯度检查点能在有限显存下训练大模型。第三是数值稳定性softmax、logsumexp、LayerNorm这些操作在数学上简单但在浮点数下需要特殊的稳定实现否则很容易出现NaN。所以从零手搓的意义不在于替代框架而在于理解框架。当你知道loss.backward()背后发生了什么你就能理解为什么有些操作会断开梯度、为什么detach()能省显存、为什么原地操作会导致反向传播报错。这些知识在排查训练问题时是救命的。3. 训练循环里那些框架帮你藏起来的细节3.1 数据加载不是简单的for循环很多人写训练循环就是for x, y in dataloader: ...觉得数据加载没什么好讲的。但我在实际项目里遇到过一个诡异的问题模型训练到第3个epoch突然变慢GPU利用率从90%掉到30%。排查了半天才发现是数据加载的num_workers设置成了0所有数据预处理都在主进程里串行执行前两个epoch数据还在内存缓存里所以快缓存用完之后就开始拖后腿。数据加载的核心矛盾是GPU计算速度远快于CPU数据预处理速度。如果数据供给跟不上GPU就会空转。解决方案是使用多进程预取让CPU在GPU计算当前batch的时候提前准备好后面的batch。PyTorch的DataLoader通过num_workers参数控制工作进程数通过prefetch_factor控制每个进程预取的batch数。但num_workers不是越大越好。我实测过一个图像分类任务在8核CPU的机器上num_workers4时GPU利用率最高调到8反而下降因为进程间通信和内存拷贝的开销超过了并行收益。经验值是num_workers CPU核心数 / 2然后根据实际GPU利用率微调。from torch.utils.data import DataLoader loader DataLoader( dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue, prefetch_factor2, persistent_workersTrue )pin_memoryTrue会把数据放到锁页内存里这样从CPU传到GPU的时候可以用DMA直接拷贝不需要经过CPU中转能省不少时间。persistent_workersTrue让工作进程在epoch之间不销毁避免每个epoch重新创建进程的开销。这两个参数在小数据集上效果不明显但在大规模训练里能带来10%以上的吞吐提升。3.2 梯度累积与混合精度小显存跑大模型的两条路显存不够是训练大模型时最常见的瓶颈。假设你有一张24GB显存的卡想训练一个参数量1B的模型光模型权重fp32就要4GB加上优化器状态Adam需要两份动量又是8GB再加上激活值很容易就爆了。这时候有两条路可以走梯度累积和混合精度。梯度累积的思路很简单本来一个batch算一次梯度更新一次参数现在改成算N个batch的梯度累加起来再更新一次。效果上等价于把batch size放大了N倍但显存占用不变。实现上只需要在反向传播前不清零梯度累积N次后再清零并更新。accumulation_steps 4 optimizer.zero_grad() for i, (x, y) in enumerate(loader): loss model(x, y) / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()注意这里loss要除以accumulation_steps否则累积后的梯度会是原来的N倍相当于把学习率放大了N倍。这个细节很多人会漏掉导致训练不稳定。混合精度则是另一条路前向和反向用fp16计算但保留一份fp32的权重副本用于参数更新。fp16的显存占用是fp32的一半计算速度在支持Tensor Core的GPU上能快2到3倍。但fp16的数值范围窄容易出现梯度下溢。解决方案是使用动态损失缩放先用一个大的缩放因子把loss放大反向传播得到放大后的梯度更新前再缩回去。如果检测到梯度出现inf或nan就减小缩放因子重试。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for x, y in loader: optimizer.zero_grad() with autocast(): loss model(x, y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()我实测下来混合精度在A100上能带来约1.8倍的训练加速显存占用降低约40%。但要注意有些操作在fp16下会不稳定比如softmax、LayerNorm、以及大数相加。这些操作需要强制用fp32计算PyTorch的autocast会自动处理大部分情况但自定义算子需要手动指定。3.3 学习率调度不是调个参数那么简单学习率是训练里最敏感的超参数没有之一。设大了loss震荡不收敛设小了训练慢到怀疑人生。但比初始学习率更重要的是调度策略什么时候降、降多少、降多快。最常见的策略是余弦退火学习率从初始值按余弦曲线降到接近零。它的好处是训练后期学习率很小模型能在损失曲面的平坦区域精细调整。但余弦退火有个问题如果训练步数设得太短学习率降得太快模型还没充分探索就进入精细调整阶段最终性能反而不好。我比较推荐的是带热重启的余弦退火每隔一定步数把学习率重置回初始值让模型有机会跳出局部最优。这个策略在图像分类任务上效果很好但在自然语言处理任务上要谨慎使用因为语言模型的损失曲面更复杂热重启可能导致训练不稳定。from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts scheduler CosineAnnealingWarmRestarts( optimizer, T_010, T_mult2, eta_min1e-6 )T_010表示第一次重启周期是10个epochT_mult2表示每次重启后周期翻倍。这样训练初期频繁重启鼓励探索后期周期变长鼓励收敛。eta_min是最小学习率不要设成0否则后期参数完全不更新。还有一个容易被忽略的点warmup。训练刚开始的时候模型参数是随机初始化的梯度方向可能很乱。如果直接用大学习率很容易把模型带偏。Warmup就是在最初几百步里把学习率从0线性增加到初始值让模型先稳定下来。这个技巧在Transformer类模型上是标配能显著提升训练稳定性。4. 推理优化训练只是上半场4.1 从训练模式到推理模式不只是eval()那么简单很多人以为推理就是model.eval()然后with torch.no_grad()完事。但实际部署的时候你会发现同样的模型推理速度可能差好几倍。差距来自哪里首先是BatchNorm和Dropout的行为不同训练时BatchNorm用当前batch的统计量推理时用全局统计量Dropout在训练时随机置零推理时关闭。model.eval()就是切换这些行为。但更深层的优化在于计算图的简化。训练时计算图需要保留中间激活值用于反向传播推理时这些都不需要。torch.no_grad()会关闭梯度追踪减少内存占用。但即使这样PyTorch的默认执行模式仍然是逐算子调度的每个算子启动一次CUDA kernelkernel启动的开销在推理时占比很高。解决方案是使用推理专用编译器比如TorchScript或ONNX Runtime。TorchScript会把模型编译成一个静态图然后进行算子融合把连续的多个算子合并成一个kernel减少启动开销。我实测过一个ResNet-50模型TorchScript推理比原生PyTorch快约1.4倍ONNX Runtime快约1.7倍。# TorchScript导出 model.eval() scripted_model torch.jit.script(model) scripted_model.save(model_scripted.pt) # ONNX导出 dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, model.onnx, opset_version13, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}} )dynamic_axes参数很重要如果不指定导出的ONNX模型会固定batch size为1部署时换个batch就得重新导出。指定之后batch维度就是动态的可以灵活调整。4.2 量化用精度换速度的边界在哪里量化是把模型的权重和激活值从fp32降到int8显存占用直接降到四分之一计算速度在支持int8指令的硬件上能快2到4倍。但量化会带来精度损失关键是找到精度和速度的平衡点。量化分两种训练后量化和量化感知训练。训练后量化最简单拿一个训练好的fp32模型用一批校准数据统计激活值的分布然后确定量化参数。但这种方法对异常值很敏感如果某个激活值特别大量化范围会被拉宽导致大部分值量化后精度很差。量化感知训练是在训练过程中模拟量化的效果让模型学会适应量化带来的精度损失。具体做法是在前向传播时插入伪量化节点把权重和激活值量化再反量化反向传播时用直通估计器把梯度直接传过去。这样训练出来的模型对量化更鲁棒。import torch.quantization as quant model.eval() model.qconfig quant.get_default_qconfig(fbgemm) quant.prepare(model, inplaceTrue) # 用校准数据跑几轮 for x, _ in calib_loader: model(x) quant.convert(model, inplaceTrue)我实测过一个BERT-base模型训练后int8量化后精度掉了约1.5个百分点量化感知训练后只掉了0.3个百分点。但量化感知训练需要额外的训练时间而且实现复杂度高不少。如果对精度要求不是特别苛刻训练后量化加校准数据精选通常就够了。注意量化后的模型在CPU上加速明显但在GPU上不一定。很多GPU对int8的支持有限量化后可能反而变慢。部署前一定要在目标硬件上实测。4.3 批处理与动态批处理吞吐量和延迟的权衡推理服务的核心指标有两个吞吐量每秒处理多少请求和延迟单个请求从进到出多长时间。这两个指标往往是矛盾的增大batch size能提高吞吐量因为GPU并行度更高但单个请求的延迟会增加因为它要等同一个batch里的其他请求一起算完。静态批处理就是攒够固定数量的请求再一起推理。实现简单但延迟不可控如果请求稀疏可能要等很久才攒够一个batch。动态批处理则是设置一个时间窗口窗口内的请求攒成一个batch窗口结束就推理。这样延迟有上界吞吐量也能接受。我一般会这样设计设置最大batch size为32最大等待时间为50毫秒。请求进来后如果batch没满就等直到攒够32个或者等了50毫秒就触发推理。这个策略在大多数场景下能平衡好吞吐和延迟。import asyncio from collections import deque class DynamicBatcher: def __init__(self, max_batch_size32, max_wait_ms50): self.max_batch_size max_batch_size self.max_wait max_wait_ms / 1000 self.queue deque() self.lock asyncio.Lock() async def add_request(self, data): async with self.lock: future asyncio.Future() self.queue.append((data, future)) if len(self.queue) self.max_batch_size: await self._process_batch() else: asyncio.create_task(self._wait_and_process()) return await future async def _wait_and_process(self): await asyncio.sleep(self.max_wait) async with self.lock: if self.queue: await self._process_batch()这段代码展示了动态批处理的核心逻辑请求进来先入队如果队列满了立即处理否则启动一个定时任务等待窗口结束后处理。实际生产环境还需要考虑超时、错误处理、优先级等但核心思想就是这样。5. 工程化落地时那些文档不会写的坑5.1 随机种子可复现性没有你想的那么简单做实验最痛苦的事情莫过于昨天跑出来一个不错的结果今天同样的代码跑出来完全不一样。罪魁祸首通常是随机种子没固定。但固定随机种子这件事比设一个seed42复杂得多。首先Python内置的random、NumPy的np.random、PyTorch的torch.random是三个独立的随机数生成器必须分别设置。其次CUDA的卷积操作默认使用非确定性算法即使种子固定多次运行结果也可能不同。需要设置torch.backends.cudnn.deterministic True和torch.backends.cudnn.benchmark False。但这样会牺牲性能因为确定性算法通常比非确定性算法慢。import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False # 对于多GPU训练 torch.use_deterministic_algorithms(True, warn_onlyTrue)warn_onlyTrue表示如果某个操作没有确定性实现只警告不报错。这在实践中很有用因为有些自定义算子确实没有确定性版本强制报错会导致训练无法进行。还有一个隐藏的坑数据加载的shuffleTrue使用全局随机状态如果num_workers0每个worker的随机状态可能不同。需要在DataLoader的worker_init_fn里为每个worker设置不同的种子否则多个epoch的数据顺序可能重复。5.2 显存泄漏那些你以为释放了其实没释放的张量PyTorch使用引用计数加垃圾回收来管理显存。当一个张量不再被引用时它的显存会被释放。但如果你不小心在某个地方保留了对中间张量的引用显存就不会释放训练几个epoch后就会OOM。最常见的泄漏场景是在训练循环里把loss或中间结果存进了列表比如losses.append(loss)。loss是一个张量它保留了整个计算图的引用导致所有中间激活值都无法释放。正确的做法是losses.append(loss.item())只存数值不存张量。另一个常见场景是在验证循环里忘记加torch.no_grad()。验证时不需要梯度但如果不加no_gradPyTorch会构建计算图占用大量显存。我见过一个项目训练时显存占用正常一到验证就OOM就是因为这个原因。排查显存泄漏可以用torch.cuda.memory_summary()查看显存分配情况或者用tracemalloc追踪Python对象的内存分配。但最有效的方法还是代码审查检查所有对张量的引用确保不需要的及时释放。5.3 分布式训练单机多卡不是简单的DataParallel单机多卡训练有两种模式DataParallel和DistributedDataParallel。DataParallel实现简单一行代码就能把模型放到多张卡上但性能很差因为它是单进程多线程主卡要负责汇总所有卡的梯度通信开销大而且Python的GIL会限制并行度。DistributedDataParallel是官方推荐的方式每个卡一个进程梯度通过NCCL通信库在GPU之间直接传输不经过CPU。但它的配置复杂得多需要初始化进程组、设置local_rank、用DistributedSampler确保每个进程读到不同的数据。import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data.distributed import DistributedSampler def setup(rank, world_size): dist.init_process_group(nccl, rankrank, world_sizeworld_size) def train(rank, world_size): setup(rank, world_size) model MyModel().to(rank) model DDP(model, device_ids[rank]) sampler DistributedSampler(dataset, num_replicasworld_size, rankrank) loader DataLoader(dataset, samplersampler, batch_size32) for epoch in range(num_epochs): sampler.set_epoch(epoch) for x, y in loader: x, y x.to(rank), y.to(rank) loss model(x, y) loss.backward() optimizer.step() optimizer.zero_grad()sampler.set_epoch(epoch)这行很容易漏掉。如果不调用每个epoch的数据划分是一样的模型会反复看到同样的数据顺序影响训练效果。这个坑我踩过当时训练loss下降正常但验证集准确率一直上不去排查了很久才发现是数据划分的问题。还有一个坑是BatchNorm在DDP下的行为。默认情况下每张卡独立计算BatchNorm的统计量这相当于减小了有效batch size。如果batch size本来就小会导致训练不稳定。解决方案是使用SyncBatchNorm它在所有卡之间同步统计量但会增加通信开销。6. 从零搭建之后我真正学到了什么走完这一整套从零搭建的流程最大的收获不是学会了某个具体技术而是建立了一种“向下追问”的思维习惯。看到loss.backward()我会想计算图是怎么建的看到model.eval()我会想哪些层的行为变了看到OOM我会想显存到底被什么占着。这种思维习惯让我在遇到新框架、新工具的时候能快速抓住本质而不是被表面的API迷惑。另一个体会是AI工程和传统软件工程最大的区别在于它的“正确性”是概率性的。传统软件要么对要么错AI系统是“大部分时候对偶尔错错多少取决于数据分布”。这意味着工程化的重点不是消除错误而是管理错误监控数据分布漂移、设置置信度阈值、准备降级方案。这些内容在任何一本深度学习教材里都不会讲但它们是AI系统真正上线后每天都要面对的问题。如果你也在走这条路我的建议是不要急着追新模型、新框架先把一个最小的系统从头到尾跑通一遍。自己写数据加载、自己写训练循环、自己写推理服务、自己处理OOM和NaN。这个过程会很痛苦但走完之后你看任何AI系统的眼光都会不一样。