拓十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
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_version=13 ) # 验证:对比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_size=32, max_wait_ms=10): 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_ratio=0.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 python=3.10 -y conda activate ai-eng # 安装核心依赖,版本锁死 pip install torch==2.1.0 --index-url https://download.pytorch.org/whl/cpu pip install onnx==1.15.0 onnxruntime==1.16.3 pip install fastapi==0.104.1 uvicorn==0.24.0 pip install numpy==1.26.2 pydantic==2.5.2

这里有个细节:onnxruntime分CPU版和GPU版,安装命令不一样。CPU版直接pip install onnxruntime,GPU版要装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_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes, opset_version=13 ) 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(title="AI推理服务") # 全局加载模型,避免每次请求都加载 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_model=InferResponse) def predict(req: InferRequest): start = time.time() try: input_array = np.array(req.data, dtype=np.float32) output = session.run(None, {input_name: input_array})[0] latency = (time.time() - start) * 1000 return InferResponse(result=output.tolist(), latency_ms=latency) except Exception as e: raise HTTPException(status_code=500, detail=str(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-smi>90%持续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系统的理解都会上一个台阶。

最后分享一个我个人的小习惯:每做完一个项目,我会写一份“事故预演”文档,列出这个系统最可能出问题的十个地方,以及对应的排查步骤。这份文档平时不看,但真出问题时能省下大量慌乱的时间。这个习惯坚持了三年,帮我扛过了好几次半夜告警。系统是人建的,就一定会出问题,区别只在于你有没有准备好。

返回列表