1. 从零手搓AI工程为什么我不建议你直接调包1.1 一个让我决定“重造轮子”的契机去年帮一个做跨境电商的朋友排查推荐系统的线上问题现象很诡异离线评估AUC涨了两个千分点全量上线后点击率反而跌了。团队里两个算法工程师查了三天最后发现是特征工程阶段一个归一化参数在训练时用了全量数据的统计量导致线上推理时分布漂移。问题本身不复杂但暴露了一个普遍现象——太多人会用sklearn的fit_transform却说不清fit到底在拟合什么。这件事之后我开始有意识地做一件事把AI工程链路里那些被高级API封装掉的环节用最原始的方式重新实现一遍。不是为了炫技而是为了在出问题时能一眼看穿是哪个环节的数学假设被打破了。这个项目标题“ai-engineering-from-scratch”吸引我的地方就在这——它代表的不是“从零学AI”而是“从零构建AI工程的基础设施”把那些黑盒拆开用最朴素的代码把数据流、梯度流、服务流串起来。这篇文章适合谁看如果你已经能用PyTorch或TensorFlow跑通模型但对数据管道、梯度计算、模型服务化这些环节的底层机制一知半解那接下来的内容会对你有用。如果你是完全的新手建议先补一下线性代数和Python基础再回来看这篇。我会从数据加载开始一路讲到模型部署每个环节都给出可运行的代码和踩坑记录。1.2 从零构建的边界哪些轮子值得造哪些不值得“从零”不等于“全部手写”。我给自己划了一条线凡是涉及数学原理和数据流转核心逻辑的环节必须手写凡是涉及工程效率和硬件适配的环节可以用成熟库。比如矩阵乘法我会用NumPy而不是自己写三重循环因为NumPy底层是BLAS手写循环在性能上差两个数量级而且这不是AI工程的核心矛盾。但反向传播的梯度计算我会手写一遍因为这是理解模型训练的关键。具体来说这个项目覆盖的模块包括数据加载与预处理管道、张量操作与自动微分、模型定义与初始化、损失函数与优化器、训练循环与验证策略、模型序列化与部署。每个模块我都会先讲清楚“为什么需要它”再给出“最小可行实现”最后补充“生产环境要注意什么”。这种结构的好处是你既能理解原理又能看到从玩具代码到生产代码之间需要补哪些东西。注意从零实现的目标是理解不是替代。生产环境该用PyTorch还是用PyTorch该用TensorRT还是用TensorRT。手写一遍是为了在出问题时知道该查哪一层。2. 数据管道从原始文件到模型可用的张量2.1 为什么数据加载值得单独设计很多人觉得数据加载就是DataLoader一包就完事了但在实际项目里数据管道往往是bug最多、性能瓶颈最明显的地方。我统计过自己经手的项目训练速度慢的原因里数据加载占四成以上。常见问题包括磁盘IO成为瓶颈、预处理在训练循环里同步执行、多进程worker之间数据重复或丢失、内存暴涨导致OOM。从零设计数据管道核心要解决三个问题怎么读IO策略、怎么变预处理逻辑、怎么送批处理与传输。我采用的设计是“索引惰性加载预取”三层结构。索引层负责维护样本ID到文件路径的映射惰性加载层在__getitem__时才真正读文件预取层用后台线程或进程提前把下一批数据准备好。这样训练循环永远不会因为等数据而空转。import os import numpy as np from torch.utils.data import Dataset, DataLoader from PIL import Image class LazyImageDataset(Dataset): def __init__(self, root_dir, transformNone): self.root_dir root_dir self.transform transform self.samples [] for label, class_name in enumerate(sorted(os.listdir(root_dir))): class_dir os.path.join(root_dir, class_name) if not os.path.isdir(class_dir): continue for fname in os.listdir(class_dir): if fname.lower().endswith((.jpg, .png, .jpeg)): self.samples.append((os.path.join(class_dir, fname), label)) def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, label这段代码看起来简单但有几个细节值得说。第一sorted(os.listdir())保证了类别顺序稳定否则每次运行标签映射可能不同导致模型评估混乱。第二convert(RGB)处理了灰度图和RGBA图避免后续通道数不一致。第三路径和标签在初始化时就全部扫描完__getitem__里只做IO和变换这样多进程worker各自持有索引副本不会重复扫描目录。2.2 预处理中的数学陷阱与实操对策预处理不只是“把图缩小、把值归一化”里面藏着不少数学假设。最常见的坑是归一化统计量的计算范围。如果你用整个数据集的均值和标准差做归一化然后划分训练集和验证集那验证集的信息已经泄露到训练过程里了。正确做法是只在训练集上计算统计量然后应用到验证集和测试集。另一个坑是数据增强的随机性控制。训练时用随机裁剪、随机翻转没问题但验证和测试时必须用确定性的中心裁剪或直接缩放。我见过有人在验证集上也开了随机增强导致同一个模型两次评估结果差好几个点排查了半天以为是模型不稳定。from torchvision import transforms # 训练集随机增强 归一化 train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.8, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # 验证集确定性预处理 val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这里的mean和std是ImageNet的统计量如果你在自己的数据集上训练应该重新计算。计算方法是遍历训练集所有图片累加每个通道的像素值和平方值最后算出均值和标准差。这个过程只需要做一次把结果存下来供后续使用。2.3 批处理与内存管理的平衡术批大小batch size的选择是个老生常谈的问题但很多人只关注“显存能塞下多大”忽略了它对训练动态的影响。批大小直接影响梯度的噪声水平批太小梯度噪声大训练震荡批太大梯度估计准但可能陷入尖锐极小值泛化变差。我的经验是先从32或64开始试如果显存允许且训练稳定再逐步翻倍同时按比例调整学习率。数据加载的另一个关键是num_workers的设置。在Linux上num_workers设为CPU核心数的70%左右比较合适比如8核机器设5到6。设太多会导致进程切换开销超过IO收益设太少则GPU等数据。Windows上多进程有额外开销建议设0先用主进程跑确认逻辑没问题再开多进程。train_loader DataLoader( train_dataset, batch_size64, shuffleTrue, num_workers6, pin_memoryTrue, # 锁页内存加速CPU到GPU传输 drop_lastTrue, # 丢弃最后不足一批的数据避免BN层报错 persistent_workersTrue # 保持worker存活避免每个epoch重新启动 )pin_memoryTrue会把数据放到锁页内存里这样从CPU传到GPU时可以用DMA直接拷贝速度更快。drop_lastTrue在批大小大于1且用了BatchNorm时很重要因为最后一批可能只有一个样本BN层在单样本上计算方差会出问题。persistent_workersTrue在PyTorch 1.7以后支持能避免每个epoch结束后worker被销毁再重建的开销。3. 自动微分把反向传播的链式法则写进代码3.1 计算图自动微分的骨架自动微分的核心是计算图。前向传播时每个操作记录下输入和输出形成一张有向无环图反向传播时从损失节点出发沿着图反向走用链式法则逐节点计算梯度。理解这一点你就能明白为什么PyTorch里loss.backward()之后要optimizer.zero_grad()——因为梯度是累加的不清零就会跨批次累积。从零实现一个简易自动微分引擎需要三个核心类Tensor包装数据和梯度、Function定义前向和反向、Graph管理节点连接。下面是一个极简版本只支持标量运算但足以展示原理。class Tensor: def __init__(self, data, requires_gradFalse): self.data float(data) self.requires_grad requires_grad self.grad 0.0 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, requires_gradself.requires_grad or other.requires_grad) def _backward(): if self.requires_grad: self.grad out.grad * 1.0 if other.requires_grad: other.grad out.grad * 1.0 out._backward _backward out._prev {self, other} return out def __mul__(self, other): other other if isinstance(other, Tensor) else Tensor(other) out Tensor(self.data * other.data, requires_gradself.requires_grad or other.requires_grad) def _backward(): if self.requires_grad: self.grad out.grad * other.data if other.requires_grad: other.grad out.grad * self.data out._backward _backward out._prev {self, other} 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 1.0 for node in reversed(topo): node._backward()这段代码里_backward函数就是链式法则的具体实现。以乘法为例out self * other那么self的梯度是out.grad * other.dataother的梯度是out.grad * self.data。backward()方法先做拓扑排序确保每个节点的梯度在它被使用之前已经计算完毕然后从输出节点开始反向调用_backward。3.2 梯度检查手写实现与框架的交叉验证手写自动微分最容易出错的地方是梯度符号和系数。我习惯用数值梯度做交叉验证对某个输入加一个极小扰动计算损失变化再除以扰动得到数值梯度然后和反向传播算出的解析梯度对比误差应该在1e-6以内。def numerical_grad(f, x, eps1e-6): 计算函数f在x处的数值梯度 return (f(x eps) - f(x - eps)) / (2 * eps) # 测试f(a,b) a*b a a Tensor(2.0, requires_gradTrue) b Tensor(3.0, requires_gradTrue) c a * b a c.backward() print(f解析梯度: da{a.grad}, db{b.grad}) # da4, db2 # 数值梯度验证 def f_a(a_val): return a_val * 3.0 a_val def f_b(b_val): return 2.0 * b_val 2.0 print(f数值梯度: da{numerical_grad(f_a, 2.0)}, db{numerical_grad(f_b, 3.0)})实测下来解析梯度和数值梯度的误差在1e-7量级说明实现正确。这个技巧在实现复杂算子比如Softmax、LayerNorm时特别有用因为它们的反向公式容易推错。3.3 从标量到张量广播机制的反向传播标量版本理解了之后扩展到张量主要多了两件事广播和归约。广播发生在形状不同的张量做运算时比如(3,1)和(1,4)相加得到(3,4)。反向传播时梯度需要沿着广播的维度求和还原到原始形状。class Tensor: # ... 前面的代码 ... def __add__(self, other): other other if isinstance(other, Tensor) else Tensor(other) out Tensor(self.data other.data, ...) def _backward(): if self.requires_grad: grad out.grad # 处理广播对多出来的维度求和 while grad.ndim self.data.ndim: grad grad.sum(axis0) for i, dim in enumerate(self.data.shape): if dim 1: grad grad.sum(axisi, keepdimsTrue) self.grad grad # other同理 out._backward _backward return out这个广播处理逻辑是先去掉多余的维度在axis0上求和再把大小为1的维度求和并保持维度。比如self形状是(3,1)out形状是(3,4)那么梯度先对axis1求和得到(3,)再在axis1上保持维度求和得到(3,1)。这个逻辑在NumPy和PyTorch里是通用的理解了就能自己实现。4. 模型训练从参数初始化到学习率调度4.1 参数初始化为什么不能全零如果把所有权重初始化为零会发生什么每个神经元的输出相同反向传播时梯度也相同更新后权重依然相同。这叫对称性问题本质上是所有神经元在学同一个东西。所以初始化必须打破对称性让不同神经元走不同的路。常用的初始化方法有Xavier和He两种。Xavier适合Sigmoid和Tanh激活函数它让每一层的输出方差保持一致He初始化适合ReLU因为ReLU把一半的输入置零了方差减半所以需要更大的初始方差来补偿。def xavier_init(fan_in, fan_out): Xavier初始化方差 2 / (fan_in fan_out) limit np.sqrt(6.0 / (fan_in fan_out)) return np.random.uniform(-limit, limit, size(fan_in, fan_out)) def he_init(fan_in, fan_out): He初始化方差 2 / fan_in std np.sqrt(2.0 / fan_in) return np.random.randn(fan_in, fan_out) * std实测中如果你用ReLU激活但用了Xavier初始化训练初期loss下降会很慢因为信号在前向传播中逐层衰减。反过来如果用Tanh但用了He初始化初期梯度会爆炸。所以初始化和激活函数要配套使用。4.2 损失函数从MSE到交叉熵的数值稳定性均方误差MSE和交叉熵Cross-Entropy是两种最常用的损失函数。MSE适合回归问题交叉熵适合分类问题。但很多人不知道的是交叉熵在实现时如果不做数值稳定处理很容易出现log(0)导致NaN。def cross_entropy_loss(logits, labels): logits: (N, C) 未经过Softmax的原始输出 labels: (N,) 整数类别标签 N logits.shape[0] # 数值稳定减去每行最大值 shifted logits - logits.max(axis1, keepdimsTrue) # LogSumExp技巧 log_sum_exp np.log(np.sum(np.exp(shifted), axis1, keepdimsTrue)) log_probs shifted - log_sum_exp # 取对应类别的log概率 correct_log_probs log_probs[np.arange(N), labels] return -np.mean(correct_log_probs)这里的关键是shifted logits - logits.max(axis1, keepdimsTrue)。减去最大值不改变Softmax的结果因为分子分母同乘一个常数但能防止exp溢出。如果不做这个处理当logits里有大于709的值时exp(709)接近float64最大值就会得到inf进而导致NaN。4.3 学习率调度从固定到余弦退火学习率是训练中最难调的参数之一。固定学习率的问题在于初期可能太大导致震荡后期可能太小导致收敛慢。我常用的策略是预热余弦退火前几个epoch从很小的学习率线性增加到初始值然后按余弦曲线衰减到接近零。def cosine_schedule(step, total_steps, warmup_steps, base_lr, min_lr1e-6): if step warmup_steps: # 线性预热 return base_lr * step / warmup_steps else: # 余弦退火 progress (step - warmup_steps) / (total_steps - warmup_steps) return min_lr 0.5 * (base_lr - min_lr) * (1 np.cos(np.pi * progress))预热的作用是让模型在初期不要因为随机初始化的大梯度而跑偏余弦退火则让模型在后期以很小的步长精细调整。实测下来这个组合比固定学习率在验证集上通常能涨1到2个点。4.4 训练循环把每个环节串起来有了前面的组件训练循环就是把它们按顺序调用。但有几个细节需要注意梯度裁剪防止爆炸、验证集评估时切换模式、保存最佳模型而不是最后一个。def train_epoch(model, loader, optimizer, criterion, clip_value1.0): model.train() total_loss 0 for batch_idx, (data, target) in enumerate(loader): optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() # 梯度裁剪防止梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), clip_value) optimizer.step() total_loss loss.item() return total_loss / len(loader) def evaluate(model, loader, criterion): model.eval() total_loss 0 correct 0 with torch.no_grad(): for data, target in loader: output model(data) total_loss criterion(output, target).item() pred output.argmax(dim1) correct (pred target).sum().item() return total_loss / len(loader), correct / len(loader.dataset)model.train()和model.eval()的切换很重要因为Dropout和BatchNorm在两种模式下的行为不同。torch.no_grad()在验证时关闭梯度计算能省显存并加速。梯度裁剪的阈值一般设1.0或5.0具体看模型深度和任务。5. 模型部署从训练脚本到可调用服务5.1 模型序列化保存什么怎么保存训练完的模型要保存下来供后续使用。PyTorch提供了两种方式保存整个模型torch.save(model)和只保存参数torch.save(model.state_dict())。我强烈建议用后者因为前者依赖模型类的定义如果代码重构了模型结构加载时会报错。# 保存 torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), val_loss: val_loss, config: model_config # 保存模型配置方便重建 }, checkpoint.pth) # 加载 checkpoint torch.load(checkpoint.pth, map_locationcpu) model MyModel(**checkpoint[config]) model.load_state_dict(checkpoint[model_state_dict])保存optimizer_state_dict是为了支持断点续训因为Adam等优化器有动量缓存不保存的话续训时会有偏差。保存config是为了在加载时能重建模型结构避免硬编码。5.2 推理优化批处理、量化与ONNX训练时追求精度推理时追求速度。最简单的优化是批处理把多个请求攒成一批一起推理能充分利用GPU并行能力。但批处理会引入延迟需要根据业务场景权衡。实时性要求高的场景用批大小1离线场景可以用32或64。更进一步是量化把float32的权重和激活值转成int8模型大小减少四分之三推理速度提升两到四倍精度损失通常在1%以内。PyTorch提供了动态量化和静态量化两种方式动态量化适合LSTM和Transformer静态量化适合CNN。# 动态量化示例 import torch.quantization quantized_model torch.quantization.quantize_dynamic( model, {torch.nn.Linear}, dtypetorch.qint8 ) torch.save(quantized_model.state_dict(), quantized_model.pth)导出ONNX是另一种常见做法它把模型转成与框架无关的格式方便用ONNX Runtime或TensorRT推理。导出时要注意指定动态维度否则批大小会被固定。dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}} )5.3 服务化用FastAPI把模型包成HTTP接口模型最终要能被业务调用最简单的方式是用FastAPI包一个HTTP接口。核心逻辑是启动时加载模型到内存收到请求时预处理输入、推理、后处理输出。from fastapi import FastAPI, UploadFile from PIL import Image import io import torch app FastAPI() model None app.on_event(startup) def load_model(): global model model MyModel() model.load_state_dict(torch.load(model.pth, map_locationcpu)) model.eval() app.post(/predict) async def predict(file: UploadFile): image Image.open(io.BytesIO(await file.read())).convert(RGB) tensor val_transform(image).unsqueeze(0) with torch.no_grad(): output model(tensor) prob torch.softmax(output, dim1) pred prob.argmax(dim1).item() return {class_id: pred, confidence: prob[0][pred].item()}这个服务用uvicorn启动uvicorn main:app --host 0.0.0.0 --port 8000。生产环境还需要考虑并发、超时、日志、监控等但核心逻辑就是这些。6. 常见问题与排查技巧实录6.1 训练不收敛从数据到梯度的排查顺序训练loss不下降是最常见的问题。我的排查顺序是先看数据再看初始化然后看学习率最后看梯度。数据问题包括标签错位、归一化参数错误、数据增强过猛。初始化问题包括全零初始化、方差过大或过小。学习率问题包括太大导致震荡、太小导致停滞。梯度问题包括梯度消失、梯度爆炸、梯度为NaN。一个快速判断方法是打印每个epoch的梯度范数。如果梯度范数在1e-7以下说明梯度消失需要检查激活函数和初始化如果在1e3以上说明梯度爆炸需要加梯度裁剪或降低学习率。6.2 过拟合与欠拟合判断标准与应对策略过拟合的标志是训练loss持续下降但验证loss开始上升欠拟合的标志是训练loss和验证loss都居高不下。过拟合的应对策略包括增加数据增强、加Dropout、加L2正则化、减小模型容量、早停。欠拟合的应对策略包括增加模型容量、减少正则化、延长训练、检查数据是否有信息量。我常用的早停策略是验证loss连续5个epoch不下降就停止训练并回滚到验证loss最低的checkpoint。这个策略简单但有效能防止过拟合浪费计算资源。6.3 显存不足从批大小到混合精度的优化路径显存不足时第一反应是减小批大小但这会影响训练动态。更好的顺序是先试混合精度训练AMP能省一半显存再试梯度累积用时间换空间最后才减小批大小。混合精度训练用float16做前向和反向float32做参数更新PyTorch里几行代码就能开启。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for data, target in loader: optimizer.zero_grad() with autocast(): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()GradScaler的作用是防止float16下梯度下溢它会自动缩放loss让梯度落在float16能表示的范围内。6.4 推理结果与训练不一致模式切换与预处理对齐训练时验证准确率90%部署后推理准确率只有60%这种问题通常出在两个地方模型没有切换到eval模式或者推理时的预处理和验证时不一致。model.eval()会关闭Dropout并让BatchNorm用移动平均如果忘了这一步Dropout会随机丢弃神经元BatchNorm会用当前批的统计量结果自然不对。预处理不一致包括验证时用了CenterCrop但推理时用了Resize、归一化参数不同、通道顺序不同RGB vs BGR。我的做法是把预处理逻辑封装成一个函数训练和推理共用避免手写两套代码。问题现象可能原因排查方法解决方案训练loss为NaN学习率过大、log(0)、除零打印每层梯度、检查损失函数降低学习率、加epsilon、梯度裁剪验证loss震荡批太小、学习率太大增大批大小、降低学习率用学习率预热、梯度累积显存OOM批太大、模型太大打印各层输出形状AMP、梯度累积、减小批推理慢未用eval模式、未量化对比训练和推理耗时eval模式、ONNX、量化准确率骤降预处理不一致、模式未切换对比训练和推理的输入统一预处理、model.eval()7. 从零构建的收获与后续扩展方向7.1 手写一遍之后我看代码的方式变了做完这个项目最大的收获不是代码本身而是调试直觉。以前看到loss不下降第一反应是“换个优化器试试”现在会先打印梯度范数判断是消失还是爆炸再决定是调初始化还是加裁剪。以前看到推理慢第一反应是“换个更小的模型”现在会先看是不是没开eval模式或者预处理在CPU上同步执行拖慢了整体。这种直觉来自对每个环节的“手感”。你知道矩阵乘法的计算量是M*N*K就知道为什么大矩阵要分块你知道Softmax的数值稳定性靠的是减去最大值就知道为什么自己实现的版本会溢出你知道BatchNorm在训练和推理时的行为不同就知道为什么忘了eval模式会导致结果不一致。这些知识在文档里都有但只有亲手写过一遍才会变成条件反射。7.2 这个框架还能往哪些方向长当前这个从零实现覆盖了AI工程的核心链路但还有几个方向可以继续扩展。一是分布式训练把数据并行和模型并行的逻辑加进来理解all_reduce和all_gather的通信模式。二是混合精度手写float16和float32的转换逻辑理解为什么需要loss scaling。三是图优化实现算子融合和内存复用理解推理引擎的优化原理。四是服务编排把模型服务、特征服务、缓存层串起来理解生产级推荐系统的架构。每个方向都可以独立成篇但核心思路是一样的把黑盒拆开用最朴素的代码把数据流和梯度流串起来。这个过程不会让你成为调包高手但会让你在调包出问题时知道该查哪一层。我在实际项目里踩过的坑大部分都能在这个从零实现的框架里找到对应的原理。理解了原理排查问题就从“猜”变成了“推”。