2026/10/3 6:06:10

从零手搓AI工程:环境搭建、数据管道与推理服务实战

从零手搓AI工程:环境搭建、数据管道与推理服务实战 1. 从零手搓AI工程为什么我不建议你直接调包很多人一听到“AI工程”这四个字第一反应就是打开某个云平台拖几个组件调几个API然后跑通一个Demo就觉得自己已经入门了。我刚开始接触这个方向的时候也是这么想的直到有一次线上推理服务在高峰期直接雪崩日志里全是显存溢出的报错我才意识到——只会调包的人永远不知道系统在什么边界上会崩。ai-engineering-from-scratch这个标题核心不在“AI”而在“from scratch”。它代表的是一种学习路径不依赖高层封装从最底层的张量运算、数据管道、模型加载、推理调度开始一层一层把AI工程的地基打牢。这条路走起来慢但走通之后你对整个系统的掌控力是完全不同的。适合谁看适合那些已经会用现成框架跑模型但一遇到性能瓶颈、内存泄漏、部署异常就束手无策的开发者也适合想真正理解AI系统内部运转机制、而不是停留在“调参侠”层面的工程师。我写这篇东西不是要否定高层工具的价值——生产环境里该用还是要用。我想说的是如果你只会在高层工具上搭积木那你的能力天花板就是那个工具的天花板。而from scratch的思路是让你自己成为那个能定义天花板的人。接下来我会从环境搭建、数据管道、模型推理、性能调优、部署运维几个维度把这条路径上的关键节点和踩坑经验完整拆开讲。2. 环境搭建别急着装CUDA先把依赖关系理清楚2.1 为什么我建议用conda而不是pip做底层环境刚入门的人最容易犯的错误就是一上来就pip install torch然后发现版本冲突、CUDA不匹配、cuDNN找不到折腾一整天环境还没跑起来。我的经验是底层AI工程环境优先用conda管理pip只用来装conda仓库里没有的包。原因很简单。PyTorch、TensorFlow这些框架和CUDA、cuDNN之间有严格的版本对应关系conda在安装时会把整个依赖链一起解析而pip只管Python包层面的依赖对系统级的CUDA库无能为力。我试过在同一台机器上用pip装不同版本的torch结果就是动态库路径互相覆盖最后只能重装系统。具体操作上我习惯这样建环境conda create -n ai-scratch python3.10 conda activate ai-scratch conda install pytorch torchvision torchaudio pytorch-cuda11.8 -c pytorch -c nvidia这里有几个细节值得说。Python版本选3.10而不是最新的3.12是因为很多底层库对3.12的支持还不完善踩坑概率高。CUDA版本选11.8而不是12.x是因为11.8的生态兼容性目前最稳大部分预编译轮子都覆盖到了。这些选择没有绝对的对错但在生产环境里稳定比新潮重要得多。2.2 验证环境是否真正可用而不是“看起来装好了”装完之后很多人跑个import torch没报错就以为万事大吉。实际上import成功只说明Python层面能找到包不代表CUDA真的能用。我习惯用下面这段代码做完整验证import torch print(torch version:, torch.__version__) print(cuda available:, torch.cuda.is_available()) print(cuda version:, torch.version.cuda) print(device count:, torch.cuda.device_count()) print(device name:, torch.cuda.get_device_name(0)) # 实际跑一次张量运算确认GPU真的在工作 x torch.randn(1000, 1000).cuda() y torch.randn(1000, 1000).cuda() z torch.mm(x, y) print(matmul result shape:, z.shape) print(z device:, z.device)重点看最后那个矩阵乘法。如果cuda.is_available()返回True但矩阵乘法报错说明CUDA运行时库有问题通常是驱动版本和CUDA版本不匹配。这种情况我遇到过好几次解决办法是查NVIDIA官方文档里的驱动-CUDA对应表把驱动升到支持目标CUDA版本的最低要求以上。提示不要用nvidia-smi显示的CUDA版本去判断能用哪个CUDA Toolkit。nvidia-smi显示的是驱动支持的最高CUDA版本实际能用哪个版本取决于你安装的CUDA Toolkit和框架编译时链接的版本。2.3 目录结构从第一天就为工程化做准备from scratch不意味着代码可以乱写。恰恰相反因为你要自己管理更多底层细节目录结构的重要性反而更高。我常用的结构是这样的ai-scratch/ ├── configs/ # 配置文件yaml格式 ├── data/ # 原始数据和处理后数据 │ ├── raw/ │ └── processed/ ├── src/ │ ├── data/ # 数据加载和预处理 │ ├── models/ # 模型定义 │ ├── engine/ # 训练和推理循环 │ └── utils/ # 工具函数 ├── scripts/ # 入口脚本 ├── checkpoints/ # 模型权重 └── logs/ # 日志这个结构的好处是当你的项目从单机脚本演进到多机部署时不需要大改代码路径。我见过太多项目一开始把所有代码堆在一个main.py里等到要拆分服务时重构成本高到让人想放弃。3. 数据管道AI工程里最容易被低估的环节3.1 为什么数据加载会成为训练瓶颈大部分人第一次做AI项目时关注点全在模型结构上觉得数据加载就是DataLoader一包就完事了。但实际跑起来会发现GPU利用率经常在30%以下徘徊大部分时间都在等数据。这就是典型的数据管道瓶颈。我做过一个实测同样的模型和硬件用默认的DataLoader配置一个epoch要45分钟把num_workers从0调到8加上pin_memoryTrue和prefetch_factor4同样的epoch降到18分钟。模型代码一行没改只是把数据供给跟上了。from torch.utils.data import DataLoader loader DataLoader( dataset, batch_size64, shuffleTrue, num_workers8, # 根据CPU核心数调整通常设为物理核心数 pin_memoryTrue, # 锁页内存加速CPU到GPU的传输 prefetch_factor4, # 每个worker预取多少batch persistent_workersTrue # 避免每个epoch重新创建worker进程 )这里每个参数都有讲究。num_workers不是越大越好超过物理核心数反而会因为进程切换开销导致性能下降。pin_memory在GPU训练时几乎必开它让数据从CPU内存到GPU显存的传输走DMA通道不占用CPU。persistent_workers在epoch数多的时候效果明显省去了反复创建销毁进程的开销。3.2 自己写Dataset时最容易踩的坑from scratch意味着你很可能需要自己实现Dataset类。这里有几个我踩过的坑每一个都让我debug了好几个小时。第一个坑是在__getitem__里做重计算。比如每次取数据都重新做一遍归一化或者resize这会让数据加载慢到无法接受。正确做法是在__init__里预处理一次或者离线预处理后存成更高效的格式。第二个坑是忽略了内存泄漏。Python的__getitem__如果引用了外部大对象而num_workers0时每个worker都会复制一份内存直接爆炸。我遇到过一次8个worker把32GB内存吃满最后发现是dataset里缓存了一个全量数据列表。第三个坑是随机性控制。如果你在__getitem__里用了random而num_workers0每个worker的随机种子可能相同导致数据增强效果退化。正确做法是在worker初始化时设置不同的种子def worker_init_fn(worker_id): np.random.seed(np.random.get_state()[1][0] worker_id) random.seed(random.randint(0, 2**32) worker_id) loader DataLoader(..., worker_init_fnworker_init_fn)3.3 数据格式选择什么时候用LMDB什么时候用WebDataset当数据量大到单机内存放不下、或者需要多机共享时就需要考虑专门的数据格式。我常用的两种方案是LMDB和WebDataset它们各有适用场景。格式适用场景优点缺点LMDB单机、小文件多、随机读取频繁读取极快、支持多进程并发写入麻烦、不适合分布式WebDataset多机、大文件、顺序读取为主天然支持分片和流式加载随机访问能力弱原始文件数据量小、调试阶段简单直观大规模时IO瓶颈严重我的经验是数据量在10万条以下、单机训练直接用原始文件加缓存就够了10万到100万条、需要频繁随机读取上LMDB超过100万条或者要分布式训练考虑WebDataset。这个分界线不是绝对的但可以作为起步参考。4. 模型推理从单条推理到批量服务的演进4.1 单条推理为什么在生产环境不可接受很多教程教你推理时都是给一张图或者一句话然后模型输出结果。这种单条推理的模式在Demo阶段没问题但放到生产环境就是灾难。原因很简单GPU的并行计算能力被完全浪费了。我做过一个对比测试同一个BERT模型单条推理的吞吐量是每秒12条batch size调到32之后吞吐量变成每秒280条。硬件没变只是把请求攒成batch一起送进GPU吞吐量提升了20多倍。这就是批处理的价值。但批处理也有代价延迟会增加。因为要等够一个batch才能送进去。所以生产环境的核心问题是如何在吞吐量和延迟之间找平衡。我的做法是设置一个最大等待时间比如10毫秒在这个时间内攒到多少算多少超时就直接送一个不满的batch。4.2 动态批处理的实现思路动态批处理dynamic batching是解决上面那个平衡问题的标准方案。核心逻辑是维护一个请求队列后台线程不断从队列里取请求攒够batch size或者达到超时时间就触发一次推理。import threading import time from queue import Queue class DynamicBatcher: def __init__(self, model, max_batch_size32, max_wait_ms10): self.model model self.max_batch_size max_batch_size self.max_wait max_wait_ms / 1000.0 self.queue Queue() self.results {} self.lock threading.Lock() self._start_worker() def _start_worker(self): def worker(): while True: batch [] start time.time() while len(batch) self.max_batch_size: timeout self.max_wait - (time.time() - start) if timeout 0: break try: item self.queue.get(timeouttimeout) batch.append(item) except Exception: break if batch: self._process(batch) t threading.Thread(targetworker, daemonTrue) t.start() def _process(self, batch): inputs [item[1] for item in batch] outputs self.model(inputs) for (req_id, _), out in zip(batch, outputs): with self.lock: self.results[req_id] out这段代码是简化版实际生产还要考虑异常处理、超时清理、结果通知等。但核心思想就是这样用队列解耦请求和推理用超时保证延迟上限。4.3 显存管理推理服务最容易崩的地方推理服务和训练最大的区别是训练时显存不够可以调小batch推理时请求是外部来的你控制不了。所以显存管理必须提前做好。我的做法是启动时先做一次显存压力测试确定在当前模型和最大batch size下峰值显存占用是多少然后留20%的余量。如果请求量超过预期宁可拒绝服务也不能让显存溢出因为一旦溢出整个服务进程就挂了所有请求都失败。import torch def estimate_peak_memory(model, max_batch_size, input_shape): torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() dummy torch.randn(max_batch_size, *input_shape).cuda() with torch.no_grad(): _ model(dummy) peak torch.cuda.max_memory_allocated() total torch.cuda.get_device_properties(0).total_memory print(fPeak: {peak/1e9:.2f}GB, Total: {total/1e9:.2f}GB) print(fUsage ratio: {peak/total*100:.1f}%) del dummy torch.cuda.empty_cache() return peak如果这个比例超过80%我就会考虑减小max_batch_size或者上模型量化、显存分页等优化手段。不要等到线上崩了才去查显存那时候损失已经造成了。5. 性能调优从能用到底层优化的几个关键手段5.1 混合精度几乎零成本的加速方案混合精度训练和推理是我最推荐的优化手段因为它改动小、收益大。核心思路是让大部分计算用FP16关键部分保留FP32既省显存又提速。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()实测下来混合精度通常能省30%到50%的显存速度提升20%到40%。但要注意有些操作在FP16下会溢出比如softmax之前的logits如果太大就会变成inf。GradScaler就是用来处理这个问题的它会动态调整loss的缩放因子。5.2 算子融合与图优化PyTorch 2.0之后引入了torch.compile可以把模型的计算图编译优化自动做算子融合。我实测过一个ResNet50torch.compile之后推理速度提升了约35%而且代码改动只有一行model torch.compile(model)但torch.compile不是万能的。动态shape、控制流复杂的模型编译后可能反而变慢因为编译本身有开销而且有些模式编译不出来会回退到eager模式。我的经验是静态shape、结构规整的模型大胆用动态shape或者包含大量Python控制流的模型先测再上。5.3 内存池与显存碎片治理长时间运行的推理服务显存碎片是个隐形杀手。表现是明明总显存够用但就是分配不出连续的大块显存导致OOM。这是因为PyTorch的缓存分配器在反复分配释放不同大小的张量后会产生碎片。解决办法有两个。一是设置环境变量PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True让分配器使用可扩展段减少碎片。二是在服务空闲时定期做一次显存整理import torch import gc def defragment_memory(): gc.collect() torch.cuda.empty_cache() torch.cuda.synchronize()我通常会在服务里加一个定时任务每处理完N个请求就调用一次。N的值根据请求频率和显存增长曲线来定我一般设在1000到5000之间。6. 部署运维让服务稳定跑下去的那些细节6.1 健康检查不是走过场很多人的健康检查就是返回一个200状态码这等于没检查。真正的健康检查应该验证模型是否加载成功、GPU是否可用、显存是否在安全范围、最近一次推理是否成功。def health_check(): status {status: ok, details: {}} # 检查GPU if not torch.cuda.is_available(): status[status] error status[details][gpu] unavailable return status # 检查显存 allocated torch.cuda.memory_allocated() total torch.cuda.get_device_properties(0).total_memory ratio allocated / total status[details][memory_ratio] f{ratio:.2%} if ratio 0.9: status[status] warning # 检查模型 try: with torch.no_grad(): dummy torch.randn(1, 3, 224, 224).cuda() _ model(dummy) status[details][model] ok except Exception as e: status[status] error status[details][model] str(e) return status这个检查会真实跑一次前向传播虽然有一点开销但能发现很多隐藏问题。我建议健康检查的频率不要太高30秒一次就够了太频繁反而影响正常推理。6.2 日志要记什么不记什么日志是排查问题的生命线但记太多会拖慢服务记太少出问题时无从下手。我的原则是请求级别记摘要异常级别记全量。正常请求只记请求ID、输入shape、推理耗时、输出shape。异常请求记完整输入、完整堆栈、当时的显存状态、GPU利用率。这样正常运行时日志量可控出问题时又有足够信息。import logging import time logger logging.getLogger(__name__) def inference_with_logging(model, inputs, req_id): start time.time() try: with torch.no_grad(): outputs model(inputs) elapsed time.time() - start logger.info(freq{req_id} shape{inputs.shape} time{elapsed:.4f}s) return outputs except Exception as e: elapsed time.time() - start mem torch.cuda.memory_allocated() / 1e9 logger.error( freq{req_id} failed after {elapsed:.4f}s fmem{mem:.2f}GB error{str(e)}, exc_infoTrue ) raise6.3 优雅退出别让正在处理的请求被腰斩服务重启时如果直接kill进程正在处理的请求全部失败。正确做法是收到退出信号后停止接受新请求等正在处理的请求完成再退出。import signal import sys shutdown_flag False def handle_signal(signum, frame): global shutdown_flag shutdown_flag True print(Shutdown signal received, draining...) signal.signal(signal.SIGTERM, handle_signal) signal.signal(signal.SIGINT, handle_signal) # 主循环里检查 while not shutdown_flag: request get_request() process(request) # 等待正在处理的请求完成 wait_for_pending() print(Shutdown complete)这个逻辑看起来简单但很多服务就是缺了这一步导致每次发布都有一批用户请求失败。优雅退出的成本很低但不做的代价很高。7. 我在这条路上踩过的几个真实坑第一个坑是过度追求新版本。有段时间我每个库都追最新版结果就是各种不兼容今天这个报错明天那个警告。后来我定了个规矩生产环境的核心依赖版本锁定只在有明确收益时才升级而且升级前必须在测试环境跑满一周。第二个坑是忽略CPU和IO的瓶颈。有次推理服务优化了半天GPU利用率还是上不去最后发现是数据预处理在CPU上成了瓶颈GPU一直在等数据。后来把预处理也搬到GPU上做或者用更高效的CPU实现问题才解决。优化之前先定位瓶颈在哪不要凭感觉瞎调。第三个坑是没有做压力测试就上线。我第一次上线推理服务时本地测试好好的线上跑了几分钟就崩了。原因是并发请求一上来显存碎片化严重加上动态批处理的队列积压最后OOM。后来我养成了习惯上线前必须用工具模拟真实并发逐步加压观察显存、延迟、吞吐量的变化曲线找到系统的拐点在哪里。第四个坑是日志级别设置不当。有次线上出问题我去查日志发现全是INFO级别的正常请求记录把关键的错误信息淹没了。后来我把正常请求的日志级别调到DEBUG只在排查特定问题时才开平时保持INFO级别只记关键事件。这些坑的共同点是它们都不是算法问题而是工程问题。而from scratch这条路的价值恰恰在于让你在早期就暴露在这些工程问题面前逼着你去理解系统的每一个环节。等你把这些都趟过一遍之后再回头看那些高层工具你会知道它们帮你做了什么也会知道它们在哪些地方可能坑你。8. 后续可以继续深入的方向如果你已经把上面这些跑通了接下来有几个方向可以继续挖。一是模型量化把FP32模型转成INT8推理速度能再提升一截但精度损失需要仔细评估。二是多卡推理用张量并行或者流水线并行把大模型拆到多张卡上这个复杂度比单卡高一个量级。三是服务网格化把预处理、推理、后处理拆成独立服务各自独立扩缩容适合请求模式差异大的场景。每个方向展开都是一个大话题但有了from scratch打下的底子你再去学这些就不会只停留在“怎么调API”的层面而是能理解背后的权衡和取舍。我个人在实际操作中的体会是底层的东西学起来慢但学会了之后上层的东西学起来快。这个投入产出比长期看是划算的。