ARTICLE DETAIL

资讯详情

深耕商务建站与企业官网运营的一线实战洞察。

ONNX模型部署实战:从导出到量化,搞定跨平台推理

ONNX模型部署实战:从导出到量化,搞定跨平台推理 模型训练好只是第一步真正让开发者头疼的是怎么把它安稳地跑起来。我自己接过不少项目几乎每个项目都会在同一关卡住GPU 上训练好的 PyTorch 模型到生产环境要么缺依赖、要么算子不兼容、要么换台机器就崩。后来我把 ONNX 这套链路用熟之后部署速度明显提了一个档次。这篇文章就从 ONNX 是什么、怎么把模型导出来、用什么跑、怎么量化优化、再到实际部署场景逐个拆开讲一遍希望能帮你绕开我踩过的那些坑。ONNXOpen Neural Network Exchange本质上是一个开放格式的模型交换标准也可以理解成深度学习框架之间的“普通话”。它不像 PyTorch、TensorFlow 那样是训练框架更像是一个把模型“打包”成标准文件的容器格式再配合 ONNX Runtime 这个跨平台推理引擎几乎能在任何环境把模型跑起来。1. 模型部署为什么难难点到底在哪1.1 训练环境和部署环境是两套东西我先说一个最常见的场景。你用 PyTorch 写完一个网络训练得很好然后尝试把它交给后端同事部署。后端环境很可能没有 GPU、没有 PyTorch甚至可能是 Windows Server 加 Java 技术栈。你总不能让对方为了跑一个模型把 Python 解释器、CUDA 依赖、torch 全家桶全装一遍吧。就算装了版本稍微对不上推理结果也是天差地别。这个问题的根源在于训练框架本身就是“重型依赖”它不是专门为推理设计的。PyTorch 的模型文件只是把权重和计算图存在一起真要执行推理还得有完整的框架运行时还得有一堆底层库支持。而生产环境恰恰追求“能跑就行、依赖越少越好”。所以模型部署的第一层难题是环境隔离。我们需要的不是训练框架而是一个很轻的、只负责执行计算图的程序。ONNX 正好切在这个位置上。1.2 ONNX 到底做了什么又是怎么做的ONNX 做的事其实很简单把模型的计算图、权重、算子版本等信息统一写到一个.onnx文件里。这个文件具有平台无关性不绑定某个特定框架也不要求目标机器上安装 PyTorch 或 TensorFlow。接下来你只要用 ONNX Runtime 去加载这个文件就能在当前设备上完成推理。打个比方训练框架是“作者”ONNX 是“标准文档格式”ONNX Runtime 是“阅读器”。作者用各自的语言写文章但导出成标准文档格式后任何阅读器都能读。这和视频行业很像不管用什么剪辑软件最终压成 MP4播放器就能放。ONNX 就是模型界的 MP4。ONNX 真正巧妙的点是它的算子规范。一个模型本质上是无数算子的组合ONNX 定义了 Conv、MatMul、Relu 这些标准的算子接口。导出模型时框架把网络翻译成一张由标准算子组成的计算图运行时读取这张图把它调度到 CPU、GPU 或者其他专用芯片上执行。这就是为什么后面讲到 NCNN、RKNN 这些端侧框架时都能从 ONNX 中转因为它们都认得这套标准接口。2. 从训练框架到 ONNX导出这一步才是关键2.1 PyTorch 导出 ONNX 的完整流程很多人以为导出就是一行命令的事实际写代码时才会发现一堆细节。PyTorch 导出 ONNX 最常用的是torch.onnx.export。我直接给一个能用的基线代码。先准备一个训练好的模型一定要切换到 eval 模式并且把 BatchNorm、Dropout 等层固定住。因为推理的时候这些层的行为和训练时完全不同。接着准备一个和训练时输入 shape 相同的 dummy input这个 dummy input 会真正跑一遍网络ONNX 导出器就是靠“跟踪”这遍计算来记录图的。import torch import torch.nn as nn class MyNet(nn.Module): def __init__(self): super().__init__() self.conv nn.Conv2d(3, 32, kernel_size3, stride2, padding1) self.fc nn.Linear(32 * 56 * 56, 10) def forward(self, x): x self.conv(x) x torch.relu(x) x x.flatten(1) return self.fc(x) model MyNet() # 训练完成后要统一加载权重并固定参数 state_dict torch.load(model.pth, map_locationcpu) model.load_state_dict(state_dict) model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version17, do_constant_foldingTrue, )导出完成后第一件事是用onnx.checker检查模型完整性再用 ONNX Runtime 原地推理一遍和 PyTorch 的推理结果对比确认误差在合理范围内。import onnx import onnxruntime as ort import numpy as np onnx.checker.check_model(model.onnx) ort_session ort.InferenceSession(model.onnx, providers[CPUExecutionProvider]) ort_inputs {ort_session.get_inputs()[0].name: dummy_input.numpy()} ort_outputs ort_session.run(None, ort_inputs)如果这一步输出和 PyTorch 的model(dummy_input).detach().numpy()数值差得很远通常是 BatchNorm 没冻结或者算子实现有差异后面会专门讲排查方法。2.2 dynamic_axes、opset 与导出陷阱导出时最容易被忽略的是动态维度。很多部署场景的输入 batch 是变化的比如你不可能要求调用方每次都按固定 batch 请求模型。如果导出时不指定dynamic_axesONNX 会直接按照 dummy input 的 shape 固定死任何其他形状的输入都会直接报错。dynamic_axes的写法支持按维度动态常见的做法是把 batch 维设成动态图片宽高也经常需要动态比如目标检测里不同大小的输入图。但要提醒一句动态维度越多运行时调度成本越高某些硬件加速器还不支持动态 shape。能固定就固定不能固定就只留 batch 维动态。opset 版本也特别重要。opset 是 ONNX 算子的版本号决定了导出器能用哪些算子和图优化。导出时如果选了很老的 opset一些新算子会不支持甚至导出直接就失败选了太新的 opset又可能面临目标运行时版本跟不上的问题。我的习惯是选 opset 12 到 17 之间ONNX Runtime 1.15 以后的版本对 17 的支持已经非常稳定了。如果你的目标平台是老版本 ONNX Runtime优先用低版本导出否则模型加载时会报 unknown opset version 一类错误。导出阶段还要注意几个陷阱。第一控制流程。PyTorch 的if x.size(0) 10:这种动态分支在 torch.onnx.export 跟踪模式下并不会真正生成分支它只是跟踪了当前输入下走的路径。第二动态 shape 下的某些操作比如torch.argmax、torch.nonzero导出的算子在不同 opset 下行为可能不一样。第三前处理和后处理最好不要放进网络里。把 resize、归一化、NMS 这些放图外用业务代码处理图越简单部署越稳。3. ONNX Runtime 就是那个“什么都能跑的引擎”3.1 安装与基础推理流程ONNX Runtime 是微软开源的高性能推理引擎也是目前 ONNX 生态里最主流的运行时。安装非常省心pip install onnxruntime # GPU 环境装这个 pip install onnxruntime-gpuCPU 版和 GPU 版是分开的包千万别两个同时装会冲突。推理代码也远比你想的简洁import onnxruntime as ort import numpy as np # providers 是可选的不传就默认按可用性排序 session ort.InferenceSession(model.onnx, providers[CPUExecutionProvider]) input_name session.get_inputs()[0].name input_shape session.get_inputs()[0].shape input_type session.get_inputs()[0].type x np.random.randn(1, 3, 224, 224).astype(np.float32) outputs session.run(None, {input_name: x}) print(outputs[0])session.run的第一个参数是输出节点名列表传None表示返回所有输出。输入必须用 numpy 数组或类似对象而且 dtype 和 shape 都必须匹配模型定义。很多新手的报错出在输入 dtype 上模型要求 float32你传了一个 float64 的数组运行时会直接抛异常。如果要看模型的输入输出结构我们用onnxruntime提供的工具或者直接用 Netron 打开.onnx文件。Netron 是查图神器部署前一定要养成“先拖进 Netron 看一眼”的习惯节点结构、输入输出名、各节点参数一目了然。3.2 执行提供程序CPU、CUDA、TensorRT、NCNN、RKNN 怎么选ONNX Runtime 支持多种执行提供程序Execution Provider简称 EP每个 EP 对应一类硬件加速方案。我在实际项目里大概是这样选型的纯 CPU 服务器用默认的 CPUExecutionProvider图优化做满线程数调对已经有不错的性能。有 NVIDIA GPU 的服务器优先考虑CUDAExecutionProvider部分场景换成TensorrtExecutionProvider有明显加速但 TensorRT 版本和算子限制比较多配置复杂度高。国产加速卡或专用硬件常见有OpenVINOExecutionProvider这个对 Intel CPU/集显非常友好。手机和嵌入式设备一般不走 ONNX Runtime而是转成 NCNN、RKNN 这类端侧格式。多个 EP 可以一起写在 providers 列表里运行时优先用第一个不可用就回退。实际开发时养成把 providers 显式传进去的习惯不要依赖默认值因为默认选的 EP 不一定是最快最稳的。session ort.InferenceSession( model.onnx, providers[ TensorrtExecutionProvider, CUDAExecutionProvider, CPUExecutionProvider ] )选 EP 时要关注每类加速器对算子和内存布局的约束。TensorRT 对动态 shape 不友好且构建 engine 阶段比较慢生产环境一般会做 engine cache。CUDA 资源占用高但通用性好。真正要上边缘设备时建议直接绕开 ONNX Runtime走 ONNX 转 NCNN 或 RKNN 的链路这在第 4 部分会展开。3.3 在 Java 场景调用 ONNX 模型最近很多人搜java onnxruntime java rmbg-2.0 人物抠图说明 Java 后端调用 ONNX 模型已经是大规模生产级的需求了。好消息是 ONNX Runtime 官方提供 Java 绑定Maven 坐标是这样的dependency groupIdcom.microsoft.onnxruntime/groupId artifactIdonnxruntime/artifactId version1.17.1/version /dependencyJava 调用 ONNX 的思路和 Python 类似创建OrtSession把输入数据塞进OnnxTensor然后跑session.run。import ai.onnxruntime.*; OrtEnvironment env OrtEnvironment.getEnvironment(); OrtSession.SessionOptions options new OrtSession.SessionOptions(); OrtSession session env.createSession(rmbg.onnx, options); long[] shape {1, 3, 640, 640}; // 假设已是归一化后的 float 数组 float[] inputData loadAndPreprocessImage(portrait.jpg); OnnxTensor inputTensor OnnxTensor.createTensor(env, FloatBuffer.wrap(inputData), shape); OrtSession.Result result session.run(java.util.Map.of(input, inputTensor)); OnnxTensor outputTensor (OnnxTensor) result.get(0).getValue(); float[][][] output (float[][][]) outputTensor.getValue();Java 绑定的体验整体还算顺滑但要注意几点输入数据必须是float[]加 shape 信息不能像 Python 那样直接塞 numpy 数组图像预处理resize、归一化、BGR/RGB 转换得在 Java 侧自己写或者用 OpenCV Java 包完成模型推理结果是OnnxTensor需要手动转成 Java 数组再做后处理。RMBG-2.0 这类人物抠图模型输出通常是一张单通道 alpha 图后处理包括 sigmoid 激活和阈值二值化。把后处理逻辑写好Java 服务就能接收图片、返回透明 PNG整个流程完全不用碰 Python。4. 导出再优化图优化、精简与 INT8 量化4.1 为什么要做图优化直接导出的 ONNX 模型不一定高效因为训练框架生成的图里常常包含冗余节点和低效融合模式。比如 BatchNorm 在推理时可以折叠到前面的 Conv 层里这种融合靠 ONNX 的图优化工具可以自动完成。常用的工具有onnxsim和onnxoptimizer。onnxsim 的简化能力更强能做一些常量折叠、冗余节点消除、shape 推断等工作。我的习惯是导出后先跑一遍 onnxsim再对比跑分大多数模型都能有一定提升有些提升还非常显著。pip install onnxsim onnxoptimizer python -m onnxsim model.onnx model_sim.onnx \ --overwrite-input-shape 1,3,224,224 \ --check-n 3--check-n 3表示简化后随机跑 3 次推理对比简化前后输出是否一致。这一步比什么都重要因为它能自动验证简化过程有没有把图改坏。注意如果模型有动态维度--overwrite-input-shape最好别写死让 onnxsim 按原始动态 shape 去推断。图优化这个环节我建议在导出后、部署前做一遍不要等到目标平台发现问题再回头改。优化后的模型在那几个典型场景里都有实测收益尤其是 CPU 场景和移动端场景图越干净越好。4.2 INT8 量化动态量化、静态量化和 QDQ量化是压缩模型、提升推理速度最狠的一招也是最近很多人搜onnx 量化 int8的原因。量化的核心思路是把 FP32 的权重和激活值用 INT8 表示精度会损失一点但内存和计算开销大幅下降。常见的三种方式按使用难度排序动态量化最简单只用校准数据做权重量化激活值在推理时实时量化。适合 LSTM、Transformer 这类权重占比大的模型实现容易精度损失也小。静态量化需要准备一组校准数据推理时统计激活值的范围然后提前做好量化参数。速度提升更明显但需要数据、需要调优。量化感知训练QAT训练时就把量化误差模拟进去精度损失最小但需要重新训练模型成本最高。ONNX Runtime 的 Python API 提供了现成接口动态量化就是两三行代码的事from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model.onnx, model_int8.onnx, weight_typeQuantType.QInt8, )静态量化需要额外指定校准数据和预处理函数代码会复杂不少但部署到 CPU 服务器或边缘设备上的收益也更值得的投资。量化后一定要跑一遍精度对比用业务里最典型的样本去测。如果精度掉得太多可以尝试只量化一部分层或者换用 QAT 流程。这里要特别注意算子的兼容性。ONNX 里不是所有算子都有 INT8 实现个别算子运行时会回退到 FP32导致收益打折甚至因为图上出现混合精度而报错。量化完用 ONNX Runtime 完整推理一遍确认没有算子缺失问题再上生产。4.3 从 ONNX 到移动端和边缘端NCNN、RKNN 的链路部署到手机、开发板这种资源受限的场景通常不会直接用 ONNX Runtime而是先转成端侧框架的格式。最近热搜里反复出现onnx转rknn int8、正点原子rk3588 部署yolov8模型整个流程、如何将模型部署到orin nano上这些都是典型的 ONNX 下游链路。NCNN 是腾讯开源的高性能神经网络推理框架在 Android、Linux 嵌入式设备上非常流行。转换工具是onnx2ncnn命令行一般长这样onnx2ncnn model.onnx model.param model.bin转换完的 NCNN 模型是param和bin两个文件一个描述结构一个存权重。NCNN 对算子支持比较全但遇到不支持的算子也会直接跳过导致转换出来的模型缺层。解决办法是回源头修改模型结构或者用 NCNN 自己支持的等价算子替换。RKNN 是瑞芯微平台的推理框架用于 RK3568、RK3588 这些板子。转换一般通过rknn-toolkit2流程是 ONNX 转 RKNN然后部署到 NPU 上跑。关键点在于调 int8 量化时需要典型的图片集做校准和 ONNX 静态量化的逻辑是一致的。在你动手转 RKNN 之前强烈建议先用 ONNX Runtime 在 PC 上把精度验证好别带着精度问题直接跨平台那样排查起来会非常痛苦。先保精度再做性能优化这是端侧部署最稳的顺序。5. 实战拆解几个典型模型部署案例5.1 OCR 场景PP-OCRv6 和 RapidOCR 的 ONNX 推理OCR 是个典型的多阶段任务。以 PP-OCR 系列为例一个完整的 OCR 链路包括文本检测、方向分类、文本识别三个模型。每个模型都可以导成 ONNX再统一部署。RapidOCR 项目其实已经帮你封装好了 ONNX 版本的 PP-OCR 模型直接下载就能用这也是很多人搜微软的onnx和rapidocr的原因。RapidOCR 内部就是同时加载了检测、分类、识别三个 ONNX 模型通过 ONNX Runtime 完成前向推理。实际部署时有几个要点。第一三个模型的输入输出 shape 不同要在代码里分别读取并做预处理。第二检测模型输出的是概率图需要做阈值化、连通域分析然后才能得到文本框坐标。第三识别模型接收的是从原始图片裁剪出来的文本区域宽度是动态的所以识别模型导出时通常要把宽度维度设成动态。如果你要在 Java 后端做 OCR我建议直接把 RapidOCR 的 Python 实现跑一遍理清每个阶段的预处理和后处理逻辑然后用pp-ocrv6 onnx java这种思路把三个 ONNX 模型串联在 Java 代码里。后处理涉及的图像旋转、透视变换用 OpenCV Java 库处理即可。5.2 人物抠图RMBG-2.0 模型与 ONNX RuntimeRMBG-2.0 是专门做人物和物品背景去除的模型输出是一张 grayscale 的 alpha 图。这个模型官方提供了 ONNX 权重下载后可以直接在 ONNX Runtime 里跑省去了自己从 PyTorch 导出的麻烦。加载模型后输入通常是1x3x1024x1024的图像需要把输入图片 resize 到这个尺寸并且做归一化。推荐的图片预处理方式是把像素值除以 255然后减均值、除方差。推理后拿到输出一般是一张1x1x1024x1024的预测图。import cv2 import numpy as np import onnxruntime as ort session ort.InferenceSession(rmbg-2.0.onnx, providers[CPUExecutionProvider]) img cv2.imread(person.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (1024, 1024), interpolationcv2.INTER_LINEAR) img img.astype(np.float32) / 255.0 # 部分版本模型要求 CHW 布局 img np.transpose(img, (2, 0, 1)) img np.expand_dims(img, axis0) result session.run(None, {session.get_inputs()[0].name: img})[0] alpha result[0, 0] # 去掉 batch 和 channel 维 alpha np.clip(alpha, 0, 1) alpha cv2.resize(alpha, (original_w, original_h))后处理时可以直接把 alpha 当成掩码和原图相乘也可以在 alpha 上再做一次高斯模糊获得更柔和的边缘。抠图模型部署的大坑在于输入尺寸和归一化方式跟训练时不匹配导致边缘发虚或整图偏暗。遇到这种情况先回源头确认模型的预处理要求别盲目调阈值。5.3 目标检测YOLOv8 的 ONNX 部署全流程YOLOv8 是目前最常用的高效检测模型之一官方已经支持直接导出 ONNX。命令行或 Python API 都行yolo export modelyolov8n.pt formatonnx dynamicTrue导出后的模型输入是1x3x640x640输出是一组随着模型版本和任务不同的 tensor。检测任务最常见的是输出形状1x84x8400或1x8400x84取决于有没有转置。这里的 84 是 4 个边界框坐标加 80 个类别置信度8400 是三个尺度下候选框的总数。部署 YOLOv8 到 ONNX Runtime 时最麻烦的是后处理里的 NMSNon-Maximum Suppression非极大值抑制。ONNX Runtime 本身不提供 NMS 算子一般要么自己写 NMS要么在导出时把 NMS 部分也放进模型里。我的建议是 NMS 放代码里做图里只保留纯前向。这样图和业务解耦排查问题也更容易。推理完后对输出按置信度过滤再在剩余框里做 IoU 阈值抑制得到最终目标框。这部分代码各项目里都有不少变体但核心就是排序、去重、贪心选择。如果你要部署到 RK3588、Orin Nano 这类板子上转换到 RKNN 或 TensorRT 之前一定先把 ONNX 模型在 PC 上跑通。板子上的算子支持更严格如果 ONNX 图里有不支持的算子转换工具虽然会报错但经常不告诉你具体是哪个节点排查起来特别费时间。我习惯用 Netron 从后往前找一层层定位是哪类算子不兼容。5.4 嵌入与重排序Dify 和 RAGFlow 里的 ONNX 模型最近dify部署rerank模型、ragflow嵌入模型部署这类热搜说明RAG检索增强生成场景里部署嵌入和重排序模型已经非常普遍。嵌入模型把文本转换成向量重排序模型对检索结果重新打分这两个模型都能用 ONNX Runtime 部署。纯 Python 环境下直接把 Hugging Face 上的模型导出成 ONNX 再推理常用工具是optimum-cli或transformers里的export模块optimum-cli export onnx --model BAAI/bge-small-zh-v1.5 bge_onnx/导出完的目录通常是一个model.onnx加一个config.json。推理时读取config.json里的输入输出名按动态 shape 构造输入得到句向量。这种部署方式特别适合 Dify、RAGFlow 这类应用和自己写的 Java/Python 服务做对接。因为在嵌入模型场景模型往往是多个微服务共用的ONNX 格式让模型可以快速迁移到不同环境不绑定某个具体框架版本。重排序模型如 bge-reranker-base的输出通常是一个相关性分数部署时也要注意输入是query和document的拼接不同模型对拼接方式的要求可能不一样。这块细节直接照搬训练时的 tokenizer 配置别自己改。6. 部署中常见的坑与排查思路6.1 输入输出与预处理不一致这是我最常遇到的一类问题而且往往在自测时发现不了一上线就被线上请求打回。比如训练时输入是 RGB部署时图片喂成了 BGR训练时归一化是除以 255部署代码却忘了除训练时输入是NCHW部署时却传了NHWC。这些差异会导致推理结果完全不对但程序又不会报错。排查思路很简单把同一个输入分别用 PyTorch 和 ONNX Runtime 跑一遍比较中间层或最终输出的差异。差异大就说明图和预处理不一致差异很小可能是浮点误差可以接受。建议把预处理和后处理写成独立函数并给每个函数加上清晰的注释。模型文件也可以命名为model_input_rgb_nchw.onnx这种带含义的名字避免部署时拿错文件。6.2 推理引擎初始化失败与 503 错误部署过程中服务启动阶段的报错是最恐怖的因为日志提示往往不够明确。最常见的是初始化失败比如Engine core initialization failed或者搜到的503 - engine core initialization failed这类问题。这一类大多跟环境依赖有关常见起因包括 CUDA 版本和 ONNX Runtime 不匹配、缺少 cuDNN 库、多个版本的 onnxruntime 冲突、GPU 驱动异常等。我的排查顺序是这样的先在纯 Python 环境下跑 ONNX Runtime 推理确认模型本身没问题。检查 CUDA、cuDNN 版本是否匹配直接用官方文档里的对照表核对。检查可执行文件和服务进程的环境变量确认动态库搜索路径正确。逐一检查onnxruntime包版本尤其避免混装onnxruntime和onnxruntime-gpu。有个很容易忽略的点是“服务端和模型的硬件架构不一致”。比如你在本机 GPU 上导出并验证的模型没有任何问题但换到只有 CPU 的机器上跑时有些算子可能没有 CPU 实现导致初始化失败。解决办法是部署前先确认目标环境能选择的提供程序不要盲目依赖默认。6.3 性能不符合预期模型能跑通但速度不达标这种情况也很常见。性能问题可以分三步排查。先看线程配置。ONNX Runtime 支持设置intra_op_num_threads和inter_op_num_threads前者控制单个算子内部并行度后者控制算子间的并行度。CPU 推理场景里线程数不是越多越好设多了反而因为线程切换导致性能下降。可以用OrtSession.SessionOptions调优一般设成物理核心数的一半左右然后二分搜索找最优值。再看 EP 选型是否合理。同一模型在不同 EP 上的性能差距可能很大。CPU 上跑得不错的模型在 GPU 上可能因为没有合适的算子实现退化成 CPU 执行性能反而下降。用session.get_providers()查看当前 session 实际用了哪些 EP。最后看是否开启了图优化。ONNX Runtime 默认开启基本优化但没有开全部。可以通过SessionOptions设置最优优化级别再结合前面的 onnxsim 简化往往还有一到两成的提升空间。写在最后的一点个人经验我做模型部署做了好几年最大的体会是先把最简单的链路跑通再谈优化。很多人一上来就想着量化、想着 TensorRT、想着转 RKNN结果模型还没在 ONNX Runtime 上验证过出了问题根本分不清是哪一步导致的。先导出、再用 CPU runtime 跑通、和原始框架对比精度这步过了再往下走整个推进过程会顺很多。另外遇到奇怪的部署问题不要急着硬调代码先把模型文件用 Netron 打开仔仔细细看一遍输入输出和中间节点。很多问题在图上就能直接看出来省下的是你写十几个调试脚本的时间。模型部署是一项“差之毫厘、谬以千里”的活图里少一个 reshape后处理多一次翻转结果都能完全走样。规范化每一步耐心做精度验证比什么技巧都管用。
返回列表
PREV
查看更多资讯
NEXT
返回资讯列表