ARTICLE DETAIL

资讯详情

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

从零手搓AI工程化流程:ONNX推理服务与批处理实战

从零手搓AI工程化流程:ONNX推理服务与批处理实战 1. 为什么我要从零手搓一套AI工程化流程第一次看到ai-engineering-from-scratch这个项目标题时我脑子里蹦出来的不是“又一个教程仓库”而是过去两年带团队踩过的那些坑。市面上讲模型原理的资料多如牛毛讲Transformer架构、讲注意力机制、讲微调技巧的内容一抓一大把但真正把“一个模型从笔记本里的demo变成线上扛住流量”的完整链路讲清楚的东西少得可怜。这个项目标题的核心价值恰恰在于“from scratch”这四个字——它不是教你调包而是带你从最底层把AI工程化的每一块砖亲手垒起来。先把话说在前头这篇内容适合三类人第一类是做算法出身、模型训得不错但一上线就抓瞎的工程师第二类是有后端或运维背景、想切入AI系统但不知道从哪下手的开发者第三类是技术负责人需要判断一套AI工程体系到底该包含哪些模块、每个模块的边界在哪。如果你只是想找个现成的API调一调那这篇内容可能不太对你的胃口因为我们要聊的是“自己造轮子”这件事。ai-engineering-from-scratch这个标题拆开来看有三个关键词AI、Engineering、From Scratch。AI是领域Engineering是方法From Scratch是态度。很多人把AI工程等同于“会跑通一个推理脚本”这是最大的误解。真正的AI工程化要解决的是数据怎么流转、模型怎么版本化、推理怎么加速、服务怎么扩缩容、效果怎么监控、线上出问题怎么回滚这一整套问题。它更像是一门“把不确定性极强的模型塞进确定性极强的生产系统”的手艺而不是单纯的算法调优。我见过太多团队模型在离线测试集上指标漂亮得不行一上线就崩。有的是因为训练和推理的特征处理逻辑不一致有的是因为线上流量分布和训练数据差了一大截还有的是因为推理服务的批处理策略没设计好GPU利用率常年趴在20%以下。这些问题没有一个是靠调模型参数能解决的全是工程问题。所以这个项目标题里的“from scratch”我理解成两层意思一是从零搭建整套工程骨架二是从原理层面理解每个组件为什么这么设计。接下来我会按照我自己实际落地过的一套流程把这件事从头到尾拆一遍。2. 整体架构设计与技术选型思路2.1 从需求反推架构先想清楚要解决什么问题动手写第一行代码之前我习惯先画一张“问题地图”。AI工程化系统要回答的核心问题其实就那么几个模型从哪来、数据怎么进、推理怎么跑、服务怎么稳、效果怎么盯。这四个问题对应到架构上就是模型管理、数据管道、推理服务、监控告警四大模块。很多人一上来就纠结用Kubernetes还是Docker Compose用Triton还是TorchServe其实顺序反了。工具是最后一步先把数据流和职责边界理清楚选型自然就清晰了。我一般会先明确几个约束条件团队规模多大、日均请求量级多少、模型更新频率如何、有没有实时性要求。这几个变量直接决定架构的复杂度。比如一个日请求量几千、模型一个月更新一次的内部工具你上全套Kubernetes加服务网格就是过度设计一台带GPU的机器跑个FastAPI加定时任务就够了。反过来如果日请求量上千万、模型每周迭代、要求P99延迟在100毫秒以内那没有一套完整的工程体系根本撑不住。提示架构设计的第一原则是“匹配当前阶段”而不是“一步到位”。我见过太多团队在日活还没过千的时候就搭了一套号称能扛百万并发的系统结果维护成本高到没人愿意碰最后反而拖慢了迭代速度。2.2 技术栈选型每个组件背后的取舍逻辑选型这件事我的原则是“成熟优先、可替换、少造轮子”。AI工程领域变化太快今天火的框架明年可能就没人维护了所以每个组件都要保证能相对容易地换掉。下面这张表是我在实际项目中反复验证过的一套组合覆盖了从数据到服务的全链路。模块选型选择理由替代方案模型训练框架PyTorch生态活跃、调试友好、动态图适合研究TensorFlow、JAX模型格式ONNX跨框架、跨硬件、推理优化工具链成熟TorchScript、SavedModel推理引擎ONNX Runtime部署简单、CPU/GPU通吃、延迟稳定TensorRT、OpenVINO服务框架FastAPI异步支持好、类型提示、自动文档Flask、Tornado容器化Docker环境一致性、部署标准化直接裸机部署编排Docker Compose起步单机够用、学习成本低Kubernetes监控Prometheus Grafana指标采集标准、可视化灵活自建日志系统模型仓库本地文件系统 版本号简单直接、无额外依赖MLflow、DVC这张表里我想重点说两个选择。第一个是ONNX作为模型中间格式。很多人觉得导出ONNX多此一举直接用PyTorch的torch.save存整个模型不就行了问题在于PyTorch的序列化格式和具体版本强绑定训练环境升级一次线上服务可能就跑不起来了。ONNX把模型的计算图固化下来和训练框架解耦推理端只需要一个ONNX Runtime就能跑环境依赖从几个G降到几百兆。第二个是Docker Compose而不是Kubernetes。对于绝大多数中小规模场景Compose的编排能力完全够用而且调试成本低得多。Kubernetes的学习曲线和运维复杂度在没有专职SRE的情况下往往是负收益。2.3 数据流设计训练和推理必须走同一条管道这是我最想强调的一点也是踩坑最多的地方。训练时的特征处理和推理时的特征处理如果走的是两套代码那线上效果和离线评估对不上几乎是必然的。我的做法是把特征处理逻辑抽成一个独立的模块训练和推理都调用同一份代码。这个模块的输入是原始数据输出是模型可以直接吃的张量中间的所有清洗、归一化、编码逻辑都封装在里面。具体实现上我会定义一个FeaturePipeline类里面包含fit和transform两个方法。fit在训练阶段调用计算并保存统计量比如均值、方差、类别编码表transform在训练和推理阶段都调用用保存好的统计量做转换。这样训练和推理的特征分布就是严格一致的。这个设计看起来简单但能省掉后面无数次的“为什么线上效果差这么多”的排查。3. 核心模块拆解与实操要点3.1 模型导出与ONNX转换的坑把PyTorch模型导出成ONNX听起来就是一行torch.onnx.export的事但实际操作中坑非常多。第一个坑是动态维度。如果你的模型支持变长输入比如文本分类里的不同长度句子导出时必须显式指定dynamic_axes否则ONNX会把输入维度固定死线上遇到不同长度的输入直接报错。第二个坑是算子不支持。PyTorch里一些自定义算子或者较新的算子ONNX可能还没有对应的实现导出时会报错或者静默失败。我的经验是导出后一定要用onnxruntime跑一遍验证对比PyTorch和ONNX的输出差异误差在1e-4以内才算通过。import torch import torch.onnx import onnxruntime as ort import numpy as np # 假设model是已经训练好的PyTorch模型 model.eval() dummy_input torch.randn(1, 3, 224, 224) # 导出ONNX注意dynamic_axes的设置 torch.onnx.export( model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size, 2: height, 3: width}, output: {0: batch_size} }, opset_version13 ) # 验证对比PyTorch和ONNX Runtime的输出 with torch.no_grad(): torch_output model(dummy_input).numpy() session ort.InferenceSession(model.onnx) onnx_output session.run(None, {input: dummy_input.numpy()})[0] diff np.abs(torch_output - onnx_output).max() print(f最大误差: {diff}) assert diff 1e-4, 导出误差过大检查算子实现注意opset_version不要盲目追新选一个ONNX Runtime稳定支持的版本就行。我一般用13或14太新的版本可能推理端还没跟上。3.2 推理服务的批处理与并发设计推理服务的性能瓶颈十有八九出在批处理策略上。单条请求单条推理GPU利用率能低到个位数纯属浪费。但批处理也不是越大越好批大小增加会带来延迟上升需要在吞吐和延迟之间找平衡点。我的做法是实现一个动态批处理队列请求进来先放进队列服务端每隔一个很短的时间窗口比如10毫秒把队列里的请求打包成一个批次送进模型推理完再拆包返回。这样既提高了GPU利用率又不会让单个请求等太久。import asyncio from collections import deque class DynamicBatcher: def __init__(self, max_batch_size32, max_wait_ms10): self.max_batch_size max_batch_size self.max_wait_ms max_wait_ms self.queue deque() self.lock asyncio.Lock() async def add_request(self, input_data): future asyncio.Future() async with self.lock: self.queue.append((input_data, future)) if len(self.queue) self.max_batch_size: await self._process_batch() return await future async def _process_batch(self): batch list(self.queue) self.queue.clear() inputs [item[0] for item in batch] futures [item[1] for item in batch] # 这里调用实际的推理函数 results await self._infer(inputs) for future, result in zip(futures, results): future.set_result(result)这个实现里有个细节值得说max_wait_ms这个参数怎么定。我的经验是先看业务能接受的P99延迟是多少然后倒推。比如业务要求P99在200毫秒以内模型单次推理耗时50毫秒那留给排队的时间最多150毫秒max_wait_ms设成10到20毫秒比较稳妥。设太大尾延迟会很难看设太小批处理效果出不来。3.3 模型版本管理与灰度发布模型更新是AI系统里最危险的操作之一。新模型上线效果变差甚至服务挂掉的情况我都遇到过。所以模型版本管理必须做扎实。我的方案是每个模型文件用模型名_版本号_时间戳的格式命名同时维护一个current软链接指向当前生效的版本。更新时先把新模型放到目录里跑一遍离线验证通过后再把current指向新版本。回滚就是把软链接指回旧版本秒级完成。灰度发布这块我一般用请求头里的用户ID做哈希按比例分流。比如新模型先接5%的流量观察一段时间指标没问题再逐步放大到100%。这个逻辑可以在服务层实现不需要动模型本身。import hashlib def route_model(user_id, new_model_ratio0.05): 根据用户ID哈希决定走新模型还是旧模型 hash_val int(hashlib.md5(str(user_id).encode()).hexdigest(), 16) if (hash_val % 100) (new_model_ratio * 100): return new_model return old_model提示灰度比例不要一次调太大我一般按5%、20%、50%、100%的节奏走每个阶段至少观察半天。指标不只看准确率还要看延迟、错误率、资源占用。4. 完整实操流程从零搭起一套可用的推理服务4.1 环境准备与依赖安装先把基础环境搭起来。我习惯用conda管理Python环境避免和系统Python打架。整个项目需要的核心依赖不多但版本要锁死不然换台机器就出问题。# 创建并激活环境 conda create -n ai-eng python3.10 -y conda activate ai-eng # 安装核心依赖版本锁死 pip install torch2.1.0 --index-url https://download.pytorch.org/whl/cpu pip install onnx1.15.0 onnxruntime1.16.3 pip install fastapi0.104.1 uvicorn0.24.0 pip install numpy1.26.2 pydantic2.5.2这里有个细节onnxruntime分CPU版和GPU版安装命令不一样。CPU版直接pip install onnxruntimeGPU版要装onnxruntime-gpu而且CUDA版本要和驱动匹配。我踩过的坑是服务器上装了CUDA 11.8但onnxruntime-gpu默认拉的是CUDA 12的版本跑起来直接报找不到动态库。解决办法是去官网查版本对应表指定安装对应CUDA版本的包。4.2 模型导出与验证脚本把训练好的模型导出成ONNX并且做严格的数值验证。这一步不能省我见过太多导出后不验证、上线才发现输出全是NaN的案例。# export_and_verify.py import torch import numpy as np import onnxruntime as ort def export_model(model, dummy_input, onnx_path, input_names, output_names, dynamic_axes): model.eval() torch.onnx.export( model, dummy_input, onnx_path, input_namesinput_names, output_namesoutput_names, dynamic_axesdynamic_axes, opset_version13 ) print(f模型已导出至 {onnx_path}) def verify_model(model, onnx_path, test_inputs): 用多组输入验证ONNX和PyTorch输出一致性 session ort.InferenceSession(onnx_path) input_name session.get_inputs()[0].name for i, test_input in enumerate(test_inputs): with torch.no_grad(): torch_out model(test_input).numpy() onnx_out session.run(None, {input_name: test_input.numpy()})[0] max_diff np.abs(torch_out - onnx_out).max() print(f测试用例 {i}: 最大误差 {max_diff:.6f}) if max_diff 1e-4: raise ValueError(f测试用例 {i} 误差过大导出可能有问题) print(所有测试用例验证通过)验证用的测试输入要覆盖各种边界情况最小batch、最大batch、全零输入、极端值输入。我一般至少准备5组确保模型在不同输入下都稳定。4.3 推理服务搭建与接口设计服务层用FastAPI接口设计要简洁清晰。核心就两个接口一个健康检查一个推理接口。推理接口的输入输出都用Pydantic模型定义这样自动生成文档前端对接也方便。# server.py from fastapi import FastAPI, HTTPException from pydantic import BaseModel import numpy as np import onnxruntime as ort import time app FastAPI(titleAI推理服务) # 全局加载模型避免每次请求都加载 session ort.InferenceSession(model.onnx) input_name session.get_inputs()[0].name class InferRequest(BaseModel): data: list # 输入数据实际项目中根据模型定义具体结构 class InferResponse(BaseModel): result: list latency_ms: float app.get(/health) def health(): return {status: ok} app.post(/predict, response_modelInferResponse) def predict(req: InferRequest): start time.time() try: input_array np.array(req.data, dtypenp.float32) output session.run(None, {input_name: input_array})[0] latency (time.time() - start) * 1000 return InferResponse(resultoutput.tolist(), latency_mslatency) except Exception as e: raise HTTPException(status_code500, detailstr(e))启动命令很简单uvicorn server:app --host 0.0.0.0 --port 8000 --workers 4。workers数量一般设成CPU核数但如果是GPU推理workers设多了反而会抢GPU资源一般设1到2个就够了靠批处理提吞吐。4.4 容器化与一键部署把整个服务打包成Docker镜像保证开发环境和线上环境一致。Dockerfile我一般写得比较精简基础镜像用官方的Python slim版减少体积。FROM python:3.10-slim WORKDIR /app # 先复制依赖文件利用Docker层缓存 COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt # 再复制代码和模型 COPY server.py . COPY model.onnx . EXPOSE 8000 CMD [uvicorn, server:app, --host, 0.0.0.0, --port, 8000, --workers, 2]构建和运行docker build -t ai-infer:1.0.0 . docker run -d --name ai-infer -p 8000:8000 --gpus all ai-infer:1.0.0--gpus all这个参数需要宿主机装好NVIDIA Container Toolkit不然容器里访问不到GPU。这个也是我踩过的坑装完驱动忘了装toolkit容器里onnxruntime一直报找不到CUDA。5. 监控、排查与性能调优实战5.1 关键指标采集与告警设置服务跑起来只是第一步能持续稳定运行才是本事。我一般会采集四类指标请求量QPS、延迟P50/P95/P99、错误率、资源占用GPU利用率、显存、CPU、内存。这四类指标用Prometheus采集Grafana做面板。告警规则我设得比较克制只对真正影响业务的情况报警错误率超过1%持续5分钟、P99延迟超过500毫秒持续5分钟、GPU显存超过90%持续10分钟。告警太多会让人麻木最后真出事了反而没人看。指标采集方式告警阈值处理动作QPS请求计数器突降50%检查上游依赖P99延迟直方图500ms持续5min检查批处理队列错误率计数器1%持续5min检查模型和输入GPU显存nvidia-smi90%持续10min降低批大小或扩容模型输出分布自定义指标均值偏移3σ检查数据漂移最后一行“模型输出分布”这个指标很多人会忽略但它特别重要。模型输出突然整体偏移往往意味着输入数据分布变了这时候即使延迟和错误率都正常业务效果可能已经崩了。我的做法是每隔一段时间采样一批输出计算均值和方差和历史基线对比偏移超过3个标准差就告警。5.2 常见问题速查与排查思路下面这张表是我实际运维中遇到频率最高的问题以及对应的排查路径。基本上照着走一遍80%的问题都能定位。现象可能原因排查步骤解决方案服务启动报CUDA错误驱动/toolkit版本不匹配检查nvidia-smi和容器内CUDA版本重装对应版本toolkit推理结果全是NaN输入未归一化或含异常值打印输入统计量加输入校验和清洗延迟突然飙升批处理队列积压查看队列长度和批大小调小max_wait_ms或扩容显存缓慢增长内存泄漏监控显存随时间变化检查是否有未释放的中间变量线上效果差于离线特征处理不一致对比训练和推理的特征输出统一特征管道ONNX输出和PyTorch不一致算子实现差异逐层对比输出换opset版本或替换算子这里重点说两个。第一个是显存缓慢增长这个问题特别隐蔽服务跑几天才崩一次。根因通常是推理过程中创建了新的张量但没有及时释放或者ONNX Runtime的某个配置导致缓存不回收。我的解决办法是在服务里加一个定时任务每隔几小时主动清理一次缓存同时用nvidia-smi持续监控一旦发现增长趋势就介入排查。第二个是线上效果差于离线这个前面提过九成是特征处理不一致导致的。排查方法很简单把线上的一条真实请求的原始数据拿下来分别走训练时的特征管道和推理时的特征管道对比输出。只要两边输出一致效果就不会差太多。5.3 性能调优的几个实用技巧调优这件事我的经验是“先测量再优化”。不要凭感觉猜瓶颈在哪用数据说话。我一般先用py-spy做火焰图看CPU时间花在哪再用nvidia-smi dmon看GPU利用率曲线。如果GPU利用率低但延迟高瓶颈大概率在数据预处理或后处理如果GPU利用率高但吞吐上不去瓶颈在批处理策略或模型本身。几个我实测有效的调优手段第一把输入预处理放到GPU上做。很多人习惯在CPU上把numpy数组处理好再传给模型但数据在CPU和GPU之间来回拷贝的开销很大。如果预处理逻辑能用GPU算子实现整体延迟能降20%到30%。第二开启ONNX Runtime的图优化。session_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL这一行能让推理速度提升10%到20%几乎零成本。第三合理设置线程数。session_options.intra_op_num_threads控制单个算子内部的并行度inter_op_num_threads控制算子之间的并行度。CPU推理时这两个参数影响很大我一般设成物理核数超线程核数反而会拖慢。注意调优参数不要一次改多个改一个测一个记录每次改动的效果。我见过有人一口气调了五六个参数结果性能提升了但不知道是哪个起的作用下次遇到类似问题还是抓瞎。6. 我在这套流程里踩过的坑和总结的经验6.1 那些文档里不会写的教训第一个教训永远不要相信“这个模型很简单不会出问题”。我接手过一个文本分类服务模型就是个简单的BERT微调离线准确率95%看起来稳得不行。结果上线第一周就出了两次事故一次是因为输入文本里混入了特殊字符导致tokenizer报错一次是因为某条请求的文本长度超过了模型最大长度限制。这两个问题在离线测试时都没暴露因为测试数据太“干净”了。后来我养成了一个习惯上线前用真实流量采样一批数据跑一遍专门找那些“脏”数据。第二个教训版本管理要从第一天就做。我早期做项目时觉得模型就一个改了就覆盖结果有一次改完发现效果变差想回滚却找不到旧版本了只能连夜重新训练。从那以后我强制要求所有模型文件必须带版本号而且旧版本至少保留最近5个。这个习惯救了我好几次。第三个教训监控要监控“业务指标”不能只监控“系统指标”。系统指标告诉你服务活着业务指标告诉你服务有没有用。我见过服务各项系统指标都正常但模型输出全是同一个值的情况因为输入特征全被某个bug置零了。如果只看系统指标这个问题可能几天都发现不了。6.2 关于“from scratch”这件事的再思考回到项目标题本身。ai-engineering-from-scratch这个“from scratch”我现在的理解比刚开始更深了一层。它不只是说从零搭系统更是说要从零理解每个组件为什么存在。你只有自己手写过批处理队列才知道为什么max_wait_ms不能设太大只有自己导出过ONNX才知道dynamic_axes为什么重要只有自己搭过监控才知道哪些指标真正有用。这些认知是调包调不出来的。当然我不是说所有东西都要自己造。生产环境里该用成熟工具就用成熟工具但前提是你理解这个工具解决了什么问题、边界在哪、什么时候会失效。这种理解力才是AI工程师和“调包侠”之间的分水岭。这套从零搭建的流程最大的价值不在于最终跑起来的那个服务而在于搭建过程中被迫想清楚的那些问题。每次我带着新人走一遍这个流程他们对AI系统的理解都会上一个台阶。最后分享一个我个人的小习惯每做完一个项目我会写一份“事故预演”文档列出这个系统最可能出问题的十个地方以及对应的排查步骤。这份文档平时不看但真出问题时能省下大量慌乱的时间。这个习惯坚持了三年帮我扛过了好几次半夜告警。系统是人建的就一定会出问题区别只在于你有没有准备好。
返回列表