从零构建AI工程能力:数据管线、训练循环与推理服务实战

发布时间:2026/9/28 14:07:08
从零构建AI工程能力:数据管线、训练循环与推理服务实战 从零构建AI工程能力这件事我前前后后折腾过好几轮。最早的时候我也觉得搞AI嘛不就是调个库、跑个模型、看指标涨没涨。后来踩的坑多了才明白真正卡住大多数人的从来不是某个算法本身而是从数据到训练到部署这一整条链路上每一环都有大量“文档里不会写、但你不做就会翻车”的细节。ai-engineering-from-scratch这个方向之所以值得认真聊是因为它逼着你把每一层都亲手搭一遍而不是站在别人的抽象之上。这篇文章适合两类人一类是刚入行、想搞清楚AI系统到底由哪些部件拼起来的工程师另一类是做了一段时间调包工作、想补上底层认知的从业者。我会按我自己实际搭过一遍的顺序把数据管线、训练循环、评估体系、推理服务这几块拆开讲重点放在“为什么这么设计”和“哪里最容易出问题”上。1. 先把“从零”的边界划清楚1.1 从零不等于从汇编开始很多人一听到“from scratch”就紧张以为要从矩阵乘法手写起。我的理解是从零指的是不依赖高层训练框架的封装而不是拒绝一切现成工具。你可以用NumPy做张量运算用PyTorch的autograd做反向传播但你要清楚每一步在算什么。真正需要你亲手实现的核心是这几块数据加载与预处理、前向传播的计算图、损失函数、反向传播的梯度计算、参数更新、评估指标、以及推理时的批处理逻辑。我一般建议的边界是这样的底层线性代数用成熟库自动微分可以用框架但训练循环、数据管线、评估逻辑必须自己写一遍。原因很简单——这三块是出问题最多的地方也是调包时你完全看不见的地方。你调model.fit()的时候数据是怎么shuffle的、padding是怎么补的、梯度是什么时候清零的全被藏起来了。一旦指标不对你连从哪查起都不知道。1.2 为什么建议用NumPy加autograd的组合纯NumPy手写反向传播当然可以但说实话对于稍微复杂一点的网络手推梯度非常容易出错而且调试成本极高。我的做法是用NumPy管数据和前向计算用PyTorch的autograd管梯度。这样你既能看到每一步的数值又不用自己推导链式法则。具体来说把输入数据转成torch.Tensor并设置requires_gradTrue前向计算用torch的算子反向直接调.backward()然后手动读取.grad做参数更新。这样做的好处是参数更新这一步完全由你控制。你可以清楚地看到学习率是怎么乘上去的、动量是怎么累积的、权重衰减是在哪一步加的。调包的时候这些都在优化器内部你只能看个大概。自己写一遍之后再回头看Adam的实现会发现它其实就是几个滑动平均加一个偏差修正没什么神秘的。1.3 环境准备里最容易忽略的三件事第一件是随机种子。我见过太多人跑出来的结果没法复现就是因为种子没固定。Python的random、NumPy的np.random、PyTorch的torch.manual_seed这三个都要设而且要在所有随机操作之前设。如果用了GPU还要设torch.cuda.manual_seed_all。第二件是数据类型的统一。NumPy默认是float64PyTorch默认是float32。你在两边来回转的时候如果不显式指定dtype很容易出现精度不一致导致的结果偏差。我的习惯是全程用float32在数据加载的时候就转好。第三件是设备管理。CPU和GPU之间的数据传输是有开销的如果你在每个batch里都来回搬数据训练速度会慢得离谱。正确的做法是在数据加载阶段就把数据放到目标设备上或者用pin memory加异步传输。这个细节在数据量大的时候影响非常明显。2. 数据管线决定上限的那一层2.1 为什么数据管线比模型结构更重要我自己的经验是同一个模型结构数据管线做得好和做得差最终指标能差出十几个点。这不是夸张。数据管线里藏着太多决策怎么划分训练验证测试集、怎么处理缺失值、怎么做归一化、怎么处理类别不平衡、怎么做数据增强。每一个决策都会直接影响模型学到什么。举个具体的例子。很多人做分类任务的时候直接用全体数据的均值和方差做归一化然后再划分训练集和验证集。这看起来没问题但实际上验证集的统计信息已经泄漏到训练过程里了。正确做法是先划分数据集然后只用训练集的统计量做归一化再应用到验证集和测试集上。这个细节在数据量小的时候影响特别大。2.2 手写Dataset和DataLoader的关键逻辑自己写数据加载核心要实现三个东西索引映射、批处理、打乱。索引映射是指你有一个样本列表每个样本怎么根据索引取出来。批处理是指怎么把多个样本拼成一个batch这里涉及到padding和mask的处理。打乱是指每个epoch怎么重新排列样本顺序。我一般会写一个简单的Dataset类实现__len__和__getitem__然后自己写一个collate函数来处理批处理。collate函数里最关键的是处理变长序列。比如文本数据每个样本长度不一样你需要pad到当前batch的最大长度同时生成一个mask矩阵标记哪些位置是padding。这个mask在后面计算损失的时候要用到不然padding的token也会贡献梯度影响模型学习。def collate_fn(batch): max_len max(len(x) for x in batch) padded np.zeros((len(batch), max_len), dtypenp.float32) mask np.zeros((len(batch), max_len), dtypenp.float32) for i, x in enumerate(batch): padded[i, :len(x)] x mask[i, :len(x)] 1.0 return padded, mask这段代码看起来简单但mask的处理是很多人会漏掉的。没有mask模型会把padding当成真实数据来学结果就是模型在短样本上表现很差。2.3 数据增强的度怎么把握数据增强是把双刃剑。用得好能显著提升泛化能力用过头会把数据本身的分布改得面目全非模型反而学不到东西。我的原则是增强操作必须保持标签不变。比如图像分类里旋转、翻转、裁剪通常不会改变类别但如果你做的是方向敏感的任务翻转就可能改变标签。另一个原则是增强的强度要跟数据量匹配。数据量少的时候可以适当加大增强力度数据量大的时候增强反而可能拖慢收敛。我一般会先不加增强跑一个baseline然后逐步加增强看验证集指标的变化。如果加了增强之后验证集指标反而下降说明增强过头了。还有一个容易被忽略的点是验证集和测试集不能做增强。这个听起来是常识但我确实见过有人在验证集上也做了随机裁剪导致每次评估结果都不一样根本没法比较。3. 训练循环把每一步都摊开看3.1 前向传播里那些容易写错的地方前向传播看起来就是一层一层算过去但实际写的时候有几个坑。第一个是维度对齐。矩阵乘法要求前一个的最后一维等于后一个的倒数第二维这个在调试的时候经常出问题。我的习惯是在每个关键步骤打印shape确认无误后再往下走。第二个是激活函数的选择。ReLU简单高效但在负半区梯度为零容易出现神经元死亡。LeakyReLU或者GELU在某些任务上表现更好但计算量稍大。我的经验是如果网络不深ReLU基本够用如果网络很深可以考虑用GELU或者加残差连接。第三个是初始化。全零初始化会让所有神经元学到同样的东西对称性无法打破。全一初始化会导致梯度爆炸或消失。常用的做法是Xavier初始化或者He初始化前者适合tanh激活后者适合ReLU激活。自己写的时候可以用np.random.randn乘以一个缩放因子来实现。3.2 损失函数不是越复杂越好我见过很多人一上来就用Focal Loss、Dice Loss这些复杂损失觉得越复杂效果越好。但实际上大多数情况下交叉熵就够用了。复杂损失函数往往是为了解决特定问题设计的比如类别极度不平衡、或者需要优化特定指标。如果你没有这些问题用复杂损失反而可能引入不必要的超参数。交叉熵的实现本身也有细节。PyTorch的CrossEntropyLoss内部已经包含了softmax所以你的模型输出应该是logits而不是概率。如果你自己先做了softmax再传给交叉熵相当于做了两次softmax结果会不对。这个坑我踩过当时指标一直上不去查了半天才发现是这里的问题。如果是多标签分类要用BCEWithLogitsLoss同样不要自己先做sigmoid。二分类的话可以用BCEWithLogitsLoss也可以用CrossEntropyLoss配合两个输出节点两者等价看个人习惯。3.3 梯度裁剪和梯度累积的实际用法梯度裁剪是防止梯度爆炸的常用手段。有两种方式按值裁剪和按范数裁剪。按值裁剪是把每个梯度元素限制在[-clip, clip]范围内按范数裁剪是当梯度向量的范数超过阈值时整体缩放。我一般用按范数裁剪因为它在所有梯度上保持方向不变只改变大小。torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)梯度累积是为了在显存不够的时候模拟大batch。做法是多次前向反向累积梯度然后再更新参数。这里的关键是每次累积前要清零梯度累积够了再更新。如果忘了清零梯度会一直累加相当于学习率被放大了很多倍。for i, batch in enumerate(loader): loss compute_loss(batch) loss.backward() if (i 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()注意最后如果剩下的batch不够一个累积周期要么丢弃要么在循环结束后再更新一次。这个细节会影响训练的一致性。3.4 学习率调度什么时候降降多少学习率是最重要的超参数没有之一。固定学习率往往不是最优的因为训练初期需要大学习率快速下降后期需要小学习率精细调整。常用的调度策略有StepLR、CosineAnnealing、ReduceLROnPlateau。我自己的习惯是用CosineAnnealing配合warmup。warmup是在训练最开始用很小的学习率逐步增加到设定值这样可以避免初期梯度不稳定导致的震荡。CosineAnnealing是让学习率按余弦曲线从最大值降到最小值平滑过渡。def get_lr(step, warmup_steps, total_steps, max_lr, min_lr): if step warmup_steps: return max_lr * step / warmup_steps progress (step - warmup_steps) / (total_steps - warmup_steps) return min_lr 0.5 * (max_lr - min_lr) * (1 math.cos(math.pi * progress))这个函数可以直接用参数根据你的训练步数来定。warmup一般占总步数的5%到10%max_lr和min_lr差一个数量级左右。4. 评估体系别被自己的指标骗了4.1 验证集和测试集到底怎么分我见过太多人把验证集和测试集混着用最后报出来的指标看着很好实际上线就崩。正确的做法是训练集用来更新参数验证集用来调超参数和选模型测试集只在最后用一次。测试集用多了你实际上是在对它过拟合。划分比例上如果数据量在十万级别以上可以按98/1/1来分。如果数据量只有几千那验证集和测试集的比例要适当加大比如70/15/15。如果数据量极少可以考虑交叉验证但交叉验证的计算成本会成倍增加。还有一个细节是划分时要保证分布一致。如果是分类任务要按类别分层抽样确保每个类别在训练验证测试里的比例一致。如果是时间序列不能随机划分要按时间顺序划分否则会用未来数据预测过去造成泄漏。4.2 指标选择准确率往往不够用准确率在类别平衡的时候还能看一旦类别不平衡就完全失效了。比如一个二分类任务正样本占1%你全预测负样本也能有99%的准确率但这个模型毫无用处。这时候要看精确率、召回率、F1分数或者AUC-ROC。精确率和召回率是一对矛盾。精确率高意味着你预测为正的样本里真正为正的比例高召回率高意味着真正为正的样本里被你找出来的比例高。具体看哪个取决于业务需求。如果是疾病筛查宁可误报也不能漏报那就看召回率。如果是垃圾邮件过滤宁可漏掉也不能误判那就看精确率。F1是精确率和召回率的调和平均适合需要平衡两者的场景。AUC-ROC衡量的是模型在不同阈值下的排序能力不受阈值选择影响适合整体评估模型质量。4.3 早停和模型保存的策略早停是为了防止过拟合。做法是每个epoch结束后在验证集上评估如果验证集指标连续N个epoch没有提升就停止训练。N一般取5到10。早停的同时要保存验证集上最好的模型参数而不是最后一个epoch的参数。我一般会保存两个东西一个是验证集指标最好的模型一个是最后一个epoch的模型。前者用于最终部署后者用于分析训练过程。保存的时候要连优化器状态一起保存这样如果训练中断可以恢复。torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_metric: best_metric, }, checkpoint.pt)这个checkpoint格式包含了恢复训练所需的所有信息。注意model.state_dict()保存的是参数的引用不是参数的副本所以保存后如果继续训练checkpoint里的参数也会变。如果要固定某个时刻的参数需要深拷贝。5. 推理服务从实验室到生产的那道坎5.1 批处理推理和单样本推理的差异训练的时候我们习惯用大batch因为并行效率高。但推理的时候请求往往是一个一个来的如果每个请求都单独跑一次前向吞吐量会非常低。解决办法是攒批把短时间内到达的多个请求攒成一个batch一起推理然后拆分结果返回。攒批的难点在于延迟和吞吐的权衡。攒的batch越大吞吐越高但每个请求的等待时间也越长。我一般会设一个最大等待时间比如10毫秒超过这个时间即使batch没满也要发出去。这样在延迟可控的前提下尽量提高吞吐。另一个差异是推理时不需要计算梯度所以要用torch.no_grad()或者torch.inference_mode()包起来。这不仅能省显存还能提速。inference_mode比no_grad更快因为它还会关闭一些版本计数相关的逻辑。5.2 模型导出和格式转换的坑训练好的模型要部署往往需要从训练框架导出成通用格式。PyTorch可以用torch.jit.trace或者torch.jit.script导出TorchScript也可以用torch.onnx.export导出ONNX。导出的时候最容易出问题的是动态维度。如果你的模型支持变长输入导出时要显式指定哪些维度是动态的。torch.onnx.export( model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch, 1: seq_len}, output: {0: batch}} )如果不指定dynamic_axes导出的模型会把batch size和序列长度固定死换个输入就报错。这个坑我在第一次导出ONNX的时候踩过当时以为导出成功就万事大吉结果上线后发现只能处理固定长度的输入。还有一个坑是算子兼容性。不是所有的PyTorch算子都能导出成ONNX有些自定义算子需要自己写映射。导出后最好用ONNX Runtime跑一遍确认输出和PyTorch一致。差异超过1e-4就要查原因。5.3 服务化部署的最小可行方案如果只是内部使用不需要搞太复杂的服务框架。我的最小方案是用FastAPI包一层加载模型暴露一个POST接口。请求进来后做预处理、推理、后处理返回JSON。这个方案简单直接适合快速验证。from fastapi import FastAPI import torch app FastAPI() model torch.jit.load(model.pt) model.eval() app.post(/predict) def predict(data: dict): tensor preprocess(data) with torch.inference_mode(): output model(tensor) return postprocess(output)如果要上生产需要考虑的东西就多了并发控制、超时处理、健康检查、日志监控、版本管理。这些每一个都可以展开讲很久但核心思路是一样的——把推理逻辑和业务逻辑解耦推理部分保持无状态方便水平扩展。5.4 性能优化的几个实用手段第一个是量化。把float32的权重转成int8模型大小能缩小到四分之一推理速度也能提升。PyTorch支持动态量化和静态量化动态量化最简单一行代码就能搞定适合LSTM和Transformer类模型。静态量化需要校准数据精度损失更小适合CNN。第二个是算子融合。把连续的多个算子合并成一个减少内存访问次数。这个在ONNX Runtime和TensorRT里都有自动优化你只需要导出模型剩下的交给推理引擎。第三个是缓存。如果某些请求的输入重复率很高可以在预处理阶段做缓存避免重复计算。比如文本分类里同一个句子多次请求预处理结果可以缓存起来。这个优化在特定场景下效果非常明显。6. 那些文档里不会写的踩坑记录6.1 损失突然变成NaN的排查链路训练过程中loss变成NaN是最常见的问题之一。我的排查顺序是这样的先看学习率是不是太大这是最常见的原因。然后看数据里有没有异常值比如inf或者极大的数。再看损失函数里有没有log(0)或者除以零的操作。最后看梯度是不是爆炸了。具体操作上我会在训练循环里加一个检查如果loss是NaN就打印当前batch的输入统计量和梯度范数然后中断训练。这样能快速定位到是哪个batch出的问题。如果是某个特定batch导致的大概率是数据里有脏数据。if torch.isnan(loss): print(fNaN loss at step {step}) print(fInput stats: min{x.min()}, max{x.max()}, mean{x.mean()}) print(fGrad norm: {torch.nn.utils.clip_grad_norm_(model.parameters(), 1e9)}) break这个检查我建议在调试阶段一直开着上线前再关掉。它带来的性能开销很小但能帮你省下大量排查时间。6.2 显存不够时的取舍策略显存不够是训练大模型时的常态。我的应对策略按优先级排序第一减小batch size这是最直接的。第二用梯度累积模拟大batch。第三用混合精度训练float16能省一半显存。第四用梯度检查点用计算时间换显存空间。第五用模型并行或者流水线并行但这会显著增加实现复杂度。混合精度训练需要注意loss scaling。float16的表示范围比float32小梯度太小时会下溢成零。解决办法是先把loss放大反向传播后再缩回来。PyTorch的amp模块自动处理了这个过程。scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): output model(input) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()这段代码是混合精度训练的标准写法。注意scaler.update()要放在optimizer.step()之后它负责动态调整缩放因子。6.3 训练速度慢的常见原因训练速度慢不一定是模型的问题很多时候是数据加载成了瓶颈。你可以用torch.utils.data.DataLoader的num_workers参数开多进程加载但要注意不是越大越好。num_workers设成CPU核心数左右比较合适太大反而会因为进程切换开销导致变慢。另一个常见原因是数据在CPU和GPU之间来回拷贝。解决办法是在DataLoader里设置pin_memoryTrue然后用tensor.to(device, non_blockingTrue)异步传输。这样数据传输和计算可以重叠速度能提升不少。还有一个容易被忽略的是同步点。如果你在训练循环里频繁调用.item()或者打印日志会强制同步CPU和GPU打断流水线。我的做法是每隔N步才记录一次日志中间的结果先累积在GPU上到记录点再一次性取回。6.4 模型不收敛时的检查清单模型不收敛的原因很多我一般按这个清单逐项检查学习率是不是太大或太小、数据有没有归一化、标签有没有对齐、损失函数选得对不对、初始化有没有问题、网络结构有没有写错、梯度有没有传过去。其中“梯度有没有传过去”是最容易被忽略的。有时候你加了一个自定义层但忘了实现反向传播或者用了detach()把梯度截断了结果就是这层参数一直不更新。检查方法是打印每层参数的梯度范数如果某一层一直是零那就有问题。for name, param in model.named_parameters(): if param.grad is not None: print(f{name}: grad_norm{param.grad.norm().item():.6f}) else: print(f{name}: grad is None)这个检查在调试新模型的时候非常有用建议养成习惯。7. 从能跑到好用之间还差什么7.1 配置管理别把参数写死在代码里我早期写代码喜欢把超参数直接写在训练脚本里改一个参数就要改代码。后来发现这样根本没法管理实验。正确的做法是把所有超参数抽到一个配置文件里用YAML或者JSON格式训练脚本只负责读取配置并执行。model: hidden_size: 256 num_layers: 4 dropout: 0.1 training: batch_size: 32 learning_rate: 0.001 epochs: 50 warmup_steps: 1000这样做的好处是每次实验只需要保存一份配置文件就能完整复现实验设置。配合git管理可以清楚地追踪每次实验改了什么。7.2 日志和可视化训练过程要看得见训练过程中如果不记录日志出了问题只能靠猜。我一般会记录这些东西每个step的loss、学习率、梯度范数每个epoch的验证集指标以及每个epoch的训练时间。这些数据可以用TensorBoard或者WandB可视化也可以简单地写进CSV文件。关键是记录的东西要能回答“为什么指标变差了”这个问题。如果只有loss曲线你只能看到loss涨了但不知道为什么。加上学习率曲线和梯度范数曲线你就能判断是学习率太大导致震荡还是梯度爆炸导致发散。7.3 代码组织从脚本到项目的演进一开始写脚本没问题一个文件从头跑到尾快速验证想法。但当实验变多之后就需要把代码拆开了。我的拆分方式是数据相关的一个模块模型定义的一个模块训练循环的一个模块评估的一个模块配置和工具函数各一个模块。每个模块只暴露必要的接口内部实现可以独立修改。这样拆的好处是换模型的时候只需要改模型模块数据管线不用动。换数据集的时候只需要改数据模块训练逻辑不用动。这种解耦在实验迭代快的时候特别重要能省下大量重复劳动。7.4 复现性让别人能跑出你的结果复现性不只是固定随机种子那么简单。你需要保证代码版本一致、依赖版本一致、数据版本一致、配置文件一致、硬件环境尽量一致。我一般会在项目里放一个requirements.txt锁定依赖版本用git管理代码用dvc或者简单的文件校验和管理数据版本。还有一个细节是CUDA的确定性。有些CUDA算子默认是非确定性的同样的输入两次运行结果可能不一样。如果对复现性要求极高可以设置torch.use_deterministic_algorithms(True)但这会牺牲一些性能而且不是所有算子都支持。8. 这套东西搭完之后我得到了什么说实话第一次完整搭完一遍之后我最大的收获不是某个具体的技术点而是对“AI系统”这四个字有了具体的感知。以前调包的时候模型就是一个黑盒输入进去输出出来中间发生了什么完全不知道。自己搭过一遍之后每一个环节的输入输出、每一个参数的物理意义、每一个操作的代价都变得清晰了。这种清晰带来的直接好处是排查问题的速度。以前指标不对我只能盲目地调学习率、换优化器、加数据。现在我能快速定位到是数据管线的问题、还是梯度计算的问题、还是评估逻辑的问题。这个能力在真实项目里比会调几个模型重要得多。另一个收获是对“简单方案”的尊重。我见过太多人一上来就搞复杂的架构、花哨的损失函数、精细的调度策略结果baseline都没跑通。实际上一个干净的数据管线加一个标准的Transformer加一个合适的训练循环就能解决大部分问题。复杂方案应该是在简单方案不够用的时候才引入的而不是一开始就上。最后说一个我自己的习惯每搭完一个模块我都会写一个最小测试用例用假数据跑一遍确认输入输出符合预期。这个习惯帮我省下了大量联调时间。比如数据管线的测试就是构造几条假样本确认batch的shape、mask的位置、标签的对齐都正确。训练循环的测试就是用一个极小的模型和极小的数据确认loss能下降。这些测试写起来很快但能在早期发现大部分低级错误。如果你也在走这条路我的建议是不要跳过任何一个环节。数据管线看起来无聊但它是地基。训练循环看起来简单但魔鬼都在细节里。评估体系看起来是走形式但它决定了你优化的方向对不对。把这些都亲手做一遍你对AI工程的理解会上一个台阶。