ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

从零搭建AI工程:数据管道、训练循环与FastAPI部署

从零搭建AI工程:数据管道、训练循环与FastAPI部署 1. 项目概述为什么我要从零开始搭一套AI工程很多人看到ai-engineering-from-scratch这个标题第一反应是又要学一个教程但真正动手做过的人会明白从零开始构建AI工程和跑通一个demo完全是两码事。我自己在经历了无数次数据集准备好了、模型训练完了、一部署就崩的循环之后决定把整套流程彻底重写一遍不依赖任何开箱即用的全家桶框架只保留最核心的底层组件然后一步步把它们组装起来。这篇文章不会教你调包而是会带你走一遍从数据管线、模型训练到服务部署的完整工程闭环。这套东西适合谁适合那些已经会用PyTorch或TensorFlow跑通MNIST但是面对真实业务数据时依然觉得无从下手的开发者也适合那些想搞清楚为什么我们的模型在实验室里99%准确率一到生产环境就拉胯的工程师。说实话从零开始搭不是为了炫技而是为了让你真正理解每一层发生了什么这样出问题时你才知道去哪里排查。整个项目的目标是不依赖任何重量级AI平台从裸机环境开始用最朴素的工具链构建一个可以处理图片分类任务、具备自动训练追踪、模型版本管理、简单API服务部署的最小可用AI工程系统。听起来好像不难但里面每一个环节都有坑我后面会逐个拆开讲。2. 整体设计思路与技术选型背后的考量2.1 为什么放弃全家桶框架一开始我也动过用现成MLOps平台的念头比如直接上Kubeflow或者MLflow全家桶但很快发现两个问题第一这些工具的学习成本甚至比AI本身还高我本来是想做一个图片分类结果花了三天配数据库和服务发现第二它们抽象层级太高出了问题你根本不知道是模型的问题还是平台的问题。所以这次彻底改变思路每个环节只选一个足够轻量、足够稳定的工具宁可多写两百行代码也要保证每个环节透明可控。具体的选型逻辑是这样的数据处理用Pandas加NumPy因为生态成熟遇到问题搜一下基本都有答案模型训练用PyTorch因为它对动态图的支持让debug更友好尤其适合我在实验阶段频繁调整网络结构日志追踪不用TensorBoard那种重量级的直接用一个SQLite加JSON文件记录每次实验的超参数和指标够用且完全可控部署服务用FastAPI加Uvicorn理由只有一个Python世界里能把并发性能和开发效率平衡到这个程度的它是最优解。2.2 目录结构一眼看清项目的每个零件从零搭建的第一步不是写代码而是先把项目目录结构设计好。一个好的结构能让你少掉一半头发。我的最终目录长这样ai-engineering-from-scratch/ ├── data/ │ ├── raw/ # 原始数据永不修改 │ ├── processed/ # 清洗后的数据 │ └── splits/ # 训练/验证/测试划分 ├── src/ │ ├── data_processing/ # 数据加载、清洗、增强 │ ├── features/ # 特征工程相关代码 │ ├── models/ # 模型定义、训练逻辑 │ ├── evaluation/ # 评估指标、误差分析 │ └── serving/ # API服务、模型加载 ├── experiments/ │ ├── runs/ # 每次实验的日志和指标 │ └── checkpoints/ # 模型权重文件 ├── configs/ # 所有可配置参数放在这里 ├── scripts/ # 一键训练、一键部署脚本 ├── requirements.txt └── README.md可能有人会觉得一个图片分类项目用这么复杂的目录有点小题大做但你想想一个真实项目迭代三个月之后如果数据预处理代码散布在三个notebook里模型版本有十几个weights文件却不知道哪个对应哪个这才是真正的灾难。这个目录结构是我踩过坑之后固定的模板你完全可以在此基础上按自己的场景增删。核心原则原始数据永远是只读的任何清洗操作生成新文件放到processed。checkpoints目录下每个权重文件必须带实验ID前缀防止同名文件互相覆盖。configs目录集中管理所有超参禁止在代码里硬编码。2.3 从零构建的三大关键环节整个项目我把它拆成了三个核心环节数据管道、训练循环、服务封装。这三个环节互相独立但又紧密衔接每一个环节都有很多细节值得展开讲。数据管道负责从原始图片路径到张量输入的完整转换包括加载、清洗、归一化、数据增强、批次采样。这里最容易犯的错误是在数据增强时顺手修改了原始图片导致模型看到的数据分布和真实分布不一致。训练循环包括前向传播、损失计算、反向传播、参数更新、学习率调度、早停逻辑、指标记录。很多人喜欢用PyTorch Lightning这种高层封装可以实现call_fit就能训练但我坚持手写训练循环因为只有手写过你才会理解loss.backward()到底对哪些张量产生了影响以及为什么某些操作要放在torch.no_grad()下面。服务封装是把训练好的模型变成一个可以响应HTTP请求的接口包含模型加载、输入校验、预测预处理、后处理、错误处理。这个环节最容易被忽略但恰恰是生产环境失败率最高的地方。3. 核心细节解析数据管道的完整实现3.1 数据格式的统一与清洗规则我们的项目场景是给一个花卉识别的数据集做分类模型数据总共有五类花雏菊、蒲公英、玫瑰、向日葵、郁金香。原始数据是从网上下载的图片压缩包没有统一的命名规则也没有标签文件——当然真实业务中经常会遇到比这更乱的情况有些图片甚至是损坏的、格式不正确的。所以第一步我写了一个扫描脚本遍历所有图片文件验证它们能否被正确解码。这一步别看简单它能筛掉大概2%的损坏文件如果不筛训练到一半突然报图片解码失败的错那才是最难受的。清洗规则很简单统一图片通道顺序全部转成RGB因为有些图片是灰度图或者带透明通道的PNG不统一的话输入维度会变。统一大小把短边缩放到84像素然后中心裁剪到64x64。为什么不直接resize成64x64因为直接拉伸会让图片比例失真影响模型学习的特征。标签编码采用字典映射{daisy: 0, dandelion: 1, rose: 2, sunflower: 3, tulip: 4}。处理完之后数据格式统一为一个宽表DataFrame列包括image_path、label、label_id然后进行训练集、验证集、测试集的划分。划分比例是7:2:1但这里要特别注意务必在划分之前对数据按类别做分层抽样否则运气不好时某一类花全部跑到训练集验证集完全没见过这种花那你训练出来的指标就是虚的。3.2 数据增强的参数选择逻辑数据增强是必须做的但参数不能照搬ImageNet那套标准。因为我们这里只有五类花、每类大概200张图总量只有1000张直接上复杂的增强策略会导致模型过拟合到增强后的伪特征上。我的选择是只做三种轻量级增强from torchvision import transforms train_transform transforms.Compose([ transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees15), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])随机水平翻转是最基础的因为花的图片不存在左右颠倒就识别不了的问题随机旋转控制在15度以内防止旋转太多把花的主体转出视野颜色抖动作用于亮度、对比度、饱和度模拟不同光照条件下的拍摄效果。数值看起来不大但组合起来就已经足够让模型见过同一个物体在不同状态下的样子。验证集和测试集绝对不做随机增强只做缩放、归一化。这条规则我强调一万次都不为过一旦验证集也做了数据增强你评估的就不是模型在新数据上的表现而是模型在增强数据集上的表现指标再好看也没有意义。3.3 自定义Dataset的注意事项使用PyTorch写自定义Dataset时最容易踩的坑是__getitem__方法里每次都重新读图片、做变换导致数据加载变成训练瓶颈。对于小数据集可能还好但一旦图片数量超过几万张不缓存就是一个灾难。我的做法是加了一个简单的内存缓存只在数据量小于5万时启用import os import torch from torch.utils.data import Dataset from PIL import Image class FlowerDataset(Dataset): def __init__(self, df, transformNone, cache_imagesFalse): self.df df.reset_index(dropTrue) self.transform transform self.cache_images cache_images self._cache {} def __len__(self): return len(self.df) def __getitem__(self, idx): row self.df.iloc[idx] img_path row[image_path] label row[label_id] if self.cache_images and img_path in self._cache: image self._cache[img_path].copy() else: image Image.open(img_path).convert(RGB) if self.cache_images: self._cache[img_path] image.copy() if self.transform: image self.transform(image) return image, torch.tensor(label, dtypetorch.long)这里有个小细节copy()是必要的因为RandomHorizontalFlip和RandomRotation会原地修改Tensor如果不copy缓存里的图片会在第一次增强后被改掉后面所有epoch拿到的都是被增强过的脏数据。这个bug非常隐蔽我当时排查了半天。4. 实操过程从零手写训练循环4.1 模型结构设计小而美的CNN既然是花卉识别分类模型的输入是3x64x64的图片输出是5个类别。我设计了一个轻量级但够用的CNN模型import torch.nn as nn class FlowerCNN(nn.Module): def __init__(self, num_classes5): super(FlowerCNN, self).__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.Flatten(), nn.Dropout(p0.3), nn.Linear(128 * 8 * 8, 256), nn.ReLU(inplaceTrue), nn.Dropout(p0.3), nn.Linear(256, num_classes), ) def forward(self, x): x self.features(x) x self.classifier(x) return x几个设计要点每个卷积层后面都接BatchNorm目的是稳定训练小型CNN没有残差连接不加BN很容易梯度异常。Dropout放在全连接层之间比例0.3就够了太大容易欠拟合。64x64经过三次MaxPool2倍尺寸变成8x8所以全连接输入维度是128*8*8这个数字不要算错算错了后面forward直接报维度错误。为什么不直接用一个预训练的ResNet因为项目初衷是from scratch预训练模型会引入外部因素干扰我们对训练流程的理解。而且1000张小数据集完全可以用小型CNN训练到不错的准确率没必要动用大模型。4.2 手写训练循环每个细节都有存在的理由我不使用PyTorch Lightning直接写原生训练循环。这一节是整个项目最核心的代码我拆开来讲。import torch import torch.nn as nn import torch.optim as optim from tqdm import tqdm def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() running_loss 0.0 correct 0 total 0 loop tqdm(train_loader, descTraining) for inputs, labels in loop: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() loop.set_postfix(lossloss.item(), acccorrect/total) epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc为什么每步都要手动optimizer.zero_grad()因为PyTorch的梯度是累积的如果不清零上一次batch的梯度会和这次的梯度加在一起这通常不是你想要的行为。loss.item()拿到的是Python数值不是Tensor用来做日志记录时不会占用计算图如果不小心把loss当成Tensor保存了下来计算图不释放内存会越涨越高。评估循环同样要写并且在torch.no_grad()和model.eval()模式下进行def evaluate(model, val_loader, criterion, device): model.eval() running_loss 0.0 correct 0 total 0 with torch.no_grad(): loop tqdm(val_loader, descValidation) for inputs, labels in loop: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) loss criterion(outputs, labels) running_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() loop.set_postfix(lossloss.item(), acccorrect/total) return running_loss / total, correct / total这两段代码是最经典的训练评估骨架不管模型换成什么只要还是监督学习分类任务这个骨架都可以复用。真正的定制化工作在于后面的超参数配置和记录机制。4.3 超参数配置与实验记录做好日志比训练本身还重要超参数我没有硬编码在Python脚本里而是放在configs目录下的一个JSON文件中例如configs/config.json{ experiment_name: exp_001_baseline, seed: 42, train_batch_size: 32, val_batch_size: 64, learning_rate: 0.001, weight_decay: 1e-4, epochs: 30, early_stop_patience: 5 }每次实验开始时我用一个ExperimentTracker类来初始化目录结构、保存配置、记录指标import json import os from datetime import datetime class ExperimentTracker: def __init__(self, configs): self.experiment_name configs[experiment_name] self.run_dir os.path.join(experiments, runs, self.experiment_name) self.checkpoint_dir os.path.join(experiments, checkpoints, self.experiment_name) os.makedirs(self.run_dir, exist_okTrue) os.makedirs(self.checkpoint_dir, exist_okTrue) with open(os.path.join(self.run_dir, config.json), w) as f: json.dump(configs, f, indent2) self.history {train_loss: [], train_acc: [], val_loss: [], val_acc: []} def log_metrics(self, epoch, train_loss, train_acc, val_loss, val_acc): self.history[train_loss].append(train_loss) self.history[train_acc].append(train_acc) self.history[val_loss].append(val_loss) self.history[val_acc].append(val_acc) record { epoch: epoch, train_loss: train_loss, train_acc: train_acc, val_loss: val_loss, val_acc: val_acc, timestamp: datetime.now().isoformat() } with open(os.path.join(self.run_dir, metrics.jsonl), a) as f: f.write(json.dumps(record) \n)为什么用JSONL而不是CSV因为JSONL每一行是一个独立的JSON对象追加日志时不需要先读整个文件再写回万一中途崩溃也不容易损坏数据。另外这种格式未来如果要把日志导入数据库或做可视化解析起来都很方便。训练过程中最关键的一个机制是早停最优模型保存。很多人训练完所有epoch后保留最后一个epoch的权重但最后一个epoch不一定是最优的可能已经过拟合了。我的做法是在每个epoch后检查验证集准确率只保存最好的那个状态best_val_acc 0.0 patience_counter 0 for epoch in range(configs[epochs]): train_loss, train_acc train_one_epoch(...) val_loss, val_acc evaluate(...) tracker.log_metrics(epoch, train_loss, train_acc, val_loss, val_acc) if val_acc best_val_acc: best_val_acc val_acc model_path os.path.join(tracker.checkpoint_dir, best_model.pth) torch.save(model.state_dict(), model_path) patience_counter 0 else: patience_counter 1 if patience_counter configs[early_stop_patience]: print(fEarly stopping at epoch {epoch}) break这里的torch.save只保存state_dict不保存整个模型对象因为加载模型时你还需要模型的类定义只保存权重文件会更轻量也避免不同版本PyTorch之间序列化格式不兼容的问题。5. 服务封装与API部署把模型变成产品5.1 模型加载的正确姿势训练完的模型如果不部署那就只是硬盘里一堆文件。部署的第一步是把best_model.pth加载回内存。这里有一个新手常犯的错误直接用torch.load拿到整个字典然后往模型里塞却忘了先实例化模型对象。正确做法import torch from models.cnn import FlowerCNN def load_model(model_path, device): model FlowerCNN(num_classes5) state_dict torch.load(model_path, map_locationdevice) model.load_state_dict(state_dict) model.to(device) model.eval() return modelmap_locationdevice是关键如果你在GPU上训练之后在CPU机器上部署不指定map_location就会报CUDA driver not available的错误。model.eval()同样重要因为某些层如BatchNorm和Dropout在训练和推理时的行为不同不切换到eval模式同一个输入每次预测的结果都会不一样。5.2 FastAPI服务实现与请求处理我选择FastAPI来搭建推理服务代码量不大但功能完整且自带Swagger文档方便测试from fastapi import FastAPI, UploadFile, File from PIL import Image import io import numpy as np import torch from torchvision import transforms app FastAPI(titleFlower Classification API) device torch.device(cuda if torch.cuda.is_available() else cpu) model load_model(experiments/checkpoints/exp_001_baseline/best_model.pth, device) class_names [daisy, dandelion, rose, sunflower, tulip] inference_transform transforms.Compose([ transforms.Resize(64), transforms.CenterCrop(64), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) app.post(/predict) async def predict(file: UploadFile File(...)): image_bytes await file.read() try: image Image.open(io.BytesIO(image_bytes)).convert(RGB) except Exception as e: return {error: f图片解码失败: {str(e)}} input_tensor inference_transform(image).unsqueeze(0).to(device) with torch.no_grad(): outputs model(input_tensor) probabilities torch.softmax(outputs, dim1) confidence, predicted_idx torch.max(probabilities, 1) return { class_name: class_names[predicted_idx.item()], confidence: round(confidence.item(), 4), probabilities: { cls: round(float(prob), 4) for cls, prob in zip(class_names, probabilities[0].tolist()) } }注意几个细节输入图片先用Resize(64)再CenterCrop(64)保证不同尺寸的输入图片都能被中心裁成固定尺寸。Resize和CenterCrop的两个64含义不同前者的64是短边缩放后的大小后者的64是最终输出尺寸。这里如果只Resize不Crop非正方形的图片会被拉伸变形导致预测效果变差。torch.softmax后的概率向量代表了模型对各类别的置信度不只是返回最高类别把完整概率分布返回给客户端有助于下游做阈值过滤比如置信度低于0.6时就返回不确定。推理代码放在torch.no_grad()下面显式告诉PyTorch这段过程不需要构建计算图能省下大量内存和计算时间。5.3 本地运行与吞吐验证服务跑起来之后我用一个小的压测脚本验证一下吞吐量。模拟发送100张真实图片到/predict接口统计平均响应时间。在我的笔记本CPU上单并发请求平均响应时间大约是35毫秒这个速度对于小型demo和内部工具完全够用。如果未来需要更高并发可以考虑加批处理机制或用多进程gunicorn启动多个Uvicorn worker但那是另一个话题了。记录一下启动服务的命令uvicorn src.serving.api:app --host 0.0.0.0 --port 8000 --reload--reload只在开发环境用生产环境必须去掉否则文件变化会自动重启进程这在线上是一个安全隐患。6. 常见问题与排查技巧实录6.1 数据读取瓶颈问题现象训练时显卡利用率很低CPU占用率却100%每个epoch耗时非常长。原因分析数据加载和模型计算没有并行。DataLoader的num_workers默认为0所有图片的读取和变换都在主进程里执行GPU空等CPU干活。解决方案把num_workers设为4或8并设置pin_memoryTrue。注意num_workers不是越大越好因为每个worker都会复制一份Dataset到独立内存中数据集很大的时候会吃爆内存。如果是Windows系统还要把训练代码放到if __name__ __main__:保护块里否则会一直报进程异常。train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue)6.2 训练loss不下降问题现象epoch跑了好几轮loss一直停留在1.6左右几乎没有变化。原因分析对于5分类问题随机初始化模型的交叉熵损失约为-ln(1/5)1.609loss一直是这个值说明模型完全没有学到任何有效特征。排查步骤检查数据标签是否真的对应图片内容把一批input和label画出来看一眼。检查输入Tensor的数值范围ToTensor会把0-255的像素值缩放到0-1区间如果忘了归一化输入过大可能导致梯度异常。检查学习率如果lr太大模型参数会震荡如果lr太小比如0.00001学习速度会极慢。我的经验是Adam优化器配0.001起步SGD则从0.01开始调。一个快速验证的trick只取一个batch训练看模型能不能把这一批的loss降到0。如果一个batch都学不会那就不是数据问题是代码问题。6.3 全连接层维度不匹配问题现象运行forward时报错mat1 and mat2 shapes cannot be multiplied。原因分析输入图片尺寸与全连接层第一层的输入维度不匹配。解决方案不要拍脑袋算维度直接打印特征提取器输出的形状dummy_input torch.randn(1, 3, 64, 64) features model.features(dummy_input) print(features.shape) # torch.Size([1, 128, 8, 8])根据打印结果把nn.Linear的第一参数改成128 * 8 * 8就完事了。这个技巧也适用于替换任意backbone之后因为你很难手算出各种池化层组合后的张量形状。6.4 推理时BatchNorm统计量不生效问题现象训练时准确率很高保存的best_model在部署后预测结果很差甚至每次都预测同一类。原因分析最经典的原因是加载模型后忘了调用model.eval()。BatchNorm在训练模式下使用每个batch的均值和方差在eval模式下才使用训练期间累积的全局统计量Dropout在训练模式下随机丢弃神经元eval模式下才保持全部激活。直接拿训练模式做推理Dropout会让每次预测结果随机化BatchNorm会依赖当前batch的统计量导致预测不稳定。解决方案在load_model函数中显式调用model.eval()这是一条绝对不能省略的代码。6.5 实验记录混乱问题现象隔了一周回来看checkpoints目录根本分不清哪个模型是哪个实验跑出来的只能逐个加载试。解决方案从根本上约束习惯——每个实验的配置、日志、checkpoint都放在以experiment_name命名的独立目录下目录名前缀带时间戳比如exp_20250120_1030_baseline_lr001。同时在best_model.pth旁边保存一个metadata.json记录准确率、训练时间、数据的hash值这样每个模型都自带身份证后续追踪会非常方便。7. 项目扩展思路与个人心得到这里一套最小的ai-engineering-from-scratch就已经完整跑通了从零构建数据管道用原生PyTorch写训练循环通过实验追踪记录每次迭代最后用FastAPI部署成可调用的推理服务。这个流程一点也不炫酷但它是所有AI工程的地基。我个人的体会是你在这一套朴素流程上积累的对每一行代码的理解会在未来处理复杂项目时成倍地回报你。如果你想继续扩展我建议按这样的优先级来第一优先级是给训练脚本加上分布式支持至少理解DistributedDataParallel的原理和用法第二优先级是引入自动化超参搜索用Optuna替换手动调参能明显提升模型效果第三优先级是把API服务封装成Docker镜像用Kubernetes部署成弹性服务这是走向生产环境的必经之路。最后再分享一个小技巧如果你打算长期维护这个项目强烈建议把所有的Python依赖版本精确记录在requirements.txt里并且加一个make format和make lint之类的脚本。AI工程首先是软件工程代码风格一致性和可读性会在你三个月后回看代码时立下汗马功劳。
返回列表