ONNX 推理优化:导出成功只是部署的第一步
一、模型能导出,不代表线上能稳定推理
PyTorch 模型导出 ONNX 后,通常可以接入 ONNX Runtime、TensorRT 或其他推理引擎。但导出成功并不等于部署完成。算子兼容性、动态 shape、数值误差、batch 策略、预处理一致性和运行时配置,都会影响线上效果。
很多部署问题不是在导出阶段暴露,而是在真实请求分布中出现。例如训练时固定长度输入,线上输入长度变化导致 shape 不匹配;PyTorch 和 ONNX Runtime 输出存在微小误差,经过后处理后变成分类差异;CPU 推理正常,GPU Provider 下某些算子性能不稳定。部署前必须做系统验证。
二、部署链路:导出、校验、压测都要保留
flowchart TD A[PyTorch 模型] --> B[ONNX 导出] B --> C[算子检查] C --> D[数值一致性校验] D --> E[推理服务封装] E --> F[性能压测] F --> G[灰度上线]ONNX 导出时应明确 opset 版本和动态轴。opset 太低可能缺少算子支持,太高可能与运行环境不兼容。动态轴可以支持可变 batch 或序列长度,但会影响优化器做静态优化。固定 shape 性能通常更好,可变 shape 灵活性更强,需要根据业务请求分布选择。
模型预处理要一并纳入版本管理。文本 tokenizer、图像 resize、归一化参数、类别映射和后处理阈值,任何一个不一致都会造成线上线下差异。只保存 ONNX 文件是不够的,部署包应包含完整推理契约。
三、数值校验:比较输出而不是只看能运行
下面示例展示 PyTorch 与 ONNX Runtime 输出的一致性检查。阈值应根据任务类型和后处理敏感度设置。
import numpy as np import onnxruntime as ort def compare_outputs(torch_output, onnx_path, inputs): session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"]) ort_inputs = {k: v.cpu().numpy() for k, v in inputs.items()} ort_output = session.run(None, ort_inputs)[0] diff = np.max(np.abs(torch_output.detach().cpu().numpy() - ort_output)) return diff对于分类任务,除了比较 logits 差异,还要比较 top1/top5 是否一致。对于生成、检测和排序任务,要比较最终业务输出。小的数值误差可能不会影响分类,但可能改变排序边界或检测框筛选结果。因此校验指标要贴近业务输出。
还应覆盖边界样本:最短输入、最长输入、空文本、异常字符、极端图像尺寸、小 batch、大 batch。只用一两个正常样本验证,很容易遗漏动态 shape 和后处理问题。
四、性能优化:运行时配置要和请求分布匹配
ONNX Runtime 提供图优化、线程数、Execution Provider 和内存 arena 等配置。CPU 场景要测试 intra-op 和 inter-op 线程数,GPU 场景要关注数据拷贝和 batch 聚合。线程数不是越高越好,过高会和服务框架线程池争抢 CPU。
batch 推理能提高吞吐,但会增加单请求等待时间。在线接口需要在吞吐和延迟之间取舍,可以采用微批处理,在短时间窗口内聚合请求。窗口过大延迟上升,窗口过小吞吐收益有限,必须通过压测确定。
上线后要监控输入长度分布、batch 分布、P95/P99 延迟、错误率、模型输出分布和资源利用率。推理优化不是一次导出,而是随着请求分布和模型版本变化持续调整。
五、总结
ONNX 推理部署要覆盖导出、算子兼容、数值一致性、预处理版本、性能压测和灰度监控。导出成功只是第一步,真正可用的部署需要证明输出正确、延迟可控、边界样本稳定。工程验证越充分,线上不确定性越低。