ONNX - 它到底快在哪又慢在哪一、ONNX 是什么ONNXOpen Neural Network Exchange本质上是一种模型的中间表示格式IR类似于编译器里的 LLVM IR。用 PyTorch 训练完一个模型之后模型的计算逻辑是用 Python 描述的跑推理的时候要经过 Python 解释器、PyTorch 的调度器、再到底层的 CUDA kernel。这条链路很长开销不小。ONNX 做的事情是把模型的计算图从 PyTorch 的世界里导出成一个独立的、与框架无关的静态计算图然后交给专门的推理引擎比如 ONNX Runtime去执行。打个比方PyTorch 训练出来的模型像一份 Python 脚本每次执行都要解释器逐行翻译ONNX 导出后的模型像一份编译好的二进制文件直接跑就行。一个典型的 ONNX 文件.onnx里面存的是计算图的拓扑结构哪些算子、怎么连接每个算子的类型和参数Conv、MatMul、Relu 等模型的权重以 protobuf 格式序列化输入输出的 shape 和数据类型这里要注意一点ONNX 本身只是一个格式规范它不负责执行。真正跑推理的是 ONNX Runtime简称 ORT或者其他兼容 ONNX 的推理引擎TensorRT、OpenVINO 等。说用 ONNX 加速实际上是用 ONNX Runtime 加速。二、ONNX Runtime 为什么能快理解了 ONNX 是什么之后关键问题来了为什么换个引擎跑就能快2.1 去掉了 Python 开销PyTorch 推理时即使模型本身的计算是在 GPU 上跑的Python 层面仍然有大量开销每个算子调用都要经过 Python 的函数调度动态图机制意味着每次 forward 都要重新构建计算图GIL全局解释器锁在多线程场景下是个瓶颈ORT 是纯 C 实现的推理引擎模型加载完之后整个推理过程不经过 Python。对于那些计算量不大但算子数量多的模型比如很多小卷积串联的网络光是去掉 Python 开销就能带来可观的提速。2.2 图优化Graph Optimization这是 ORT 最核心的加速手段。静态计算图的好处是引擎可以在推理之前对整张图做全局优化。常见的优化包括算子融合Operator Fusion把多个相邻算子合并成一个。比如 Conv BatchNorm ReLU 三个算子在训练时是分开执行的因为 BatchNorm 在训练和推理时行为不同但在推理时可以融合成一个 kernel。这样做的好处是减少了 kernel launch 的次数也减少了中间结果在显存里的读写。常量折叠Constant Folding如果计算图里有些节点的输入全是常量比如某个 reshape 的 shape 参数那就在加载时直接算好推理时跳过。冗余节点消除去掉那些对输出没有影响的计算节点比如连续的两次 transpose 如果互为逆操作就直接删掉。内存规划Memory Planning分析整张图的数据流提前规划好每个中间张量的内存分配和复用策略避免运行时频繁的 malloc/free。2.3 硬件特化的执行后端ORT 支持多种 Execution ProviderEPEP适用硬件特点CPU EP通用 CPU使用 MLAS 数学库针对 x86/ARM 做了 SIMD 优化CUDA EPNVIDIA GPU调用 cuDNN/cuBLASTensorRT EPNVIDIA GPU进一步将子图编译为 TensorRT engineDirectML EPWindows GPU通过 DirectX 12 调用 GPU支持 AMD/Intel/NVIDIAOpenVINO EPIntel CPU/GPU/VPUIntel 硬件上的深度优化选对 EP 很重要。同一个模型用 CPU EP 和用 CUDA EP 的性能差距可以是数量级的。三、什么时候 ONNX 能加速说了这么多原理落到实际场景里以下情况用 ONNX 通常能获得明显加速3.1 模型结构固定、输入 shape 固定这是 ONNX 最舒服的场景。输入 shape 固定意味着 ORT 可以在加载时就把所有优化做到位包括内存预分配、kernel 选择、算子融合等。典型例子固定长度的音频帧处理比如每次处理 1024 个采样点固定分辨率的图像分类固定序列长度的 NLP 推理3.2 算子数量多、单个算子计算量小前面说了Python 调度开销是按算子数量线性增长的。如果一个模型有几百个小算子比如 MobileNet 这类轻量网络PyTorch 推理时大量时间花在调度上而不是计算上。换成 ORT 之后调度开销几乎为零提速非常明显。实测数据参考一个包含 200 算子的轻量语音增强模型PyTorch 推理耗时 12ms转 ONNX 后 ORT CPU 推理耗时 3ms快了 4 倍。3.3 需要部署到非 Python 环境如果最终要把模型部署到 C 服务、移动端、嵌入式设备上ONNX 几乎是必经之路。ORT 提供 C/C、C#、Java、JavaScript 等语言的 API不依赖 Python 环境。3.4 CPU 推理场景在纯 CPU 推理的场景下ORT 的 MLAS 库针对矩阵运算做了大量 SIMD 优化AVX2、AVX-512、NEON 等通常比 PyTorch 的 CPU 后端基于 MKL 或 OpenBLAS更快尤其是在 batch size 较小的时候。四、什么时候 ONNX 反而更慢这才是很多人踩坑的地方。4.1 动态 shape 场景如果模型的输入 shape 每次都不一样比如变长序列、不同分辨率的图像ORT 的很多优化就失效了无法预分配内存每次推理都要重新规划算子融合的某些模式依赖固定 shape动态 shape 下无法触发某些 EP比如 TensorRT需要为每个新 shape 重新编译 engine第一次遇到新 shape 时会有巨大的延迟PyTorch 的动态图机制天然支持动态 shape反而没有这个问题。实际建议如果必须处理变长输入可以用 padding 把输入对齐到几个固定的 shape 档位比如 256、512、1024然后为每个档位各导出一个 ONNX 模型或者使用 ORT 的动态 shape 支持但接受一定的性能损失。4.2 模型包含不支持的算子ONNX 的算子集opset是有限的。如果模型里用了某些 PyTorch 独有的算子或者自定义算子导出时可能会导出失败被拆解成一堆等价但低效的基础算子组合后者尤其隐蔽。表面上导出成功了但实际推理速度反而变慢了因为原本一个高效的自定义 kernel 被替换成了十几个基础算子的串联。遇到这种情况要么给 ORT 注册自定义算子Custom Operator要么考虑换用其他推理引擎。4.3 模型本身计算量很大瓶颈在 GPU kernel如果模型的瓶颈是大矩阵乘法或大卷积比如 ResNet-152、大型 Transformer那 Python 调度开销相对于 GPU 计算时间来说微不足道。这种情况下PyTorch 和 ORT 调用的底层 kernel 是一样的都是 cuDNN/cuBLAS性能差距很小甚至可能因为 ORT 的 kernel 选择策略不如 PyTorch 的 autotuner 而略慢。简单说模型越大、单次计算越重ONNX 的加速比越小。4.4 导出过程引入了额外开销PyTorch 导出 ONNX 时使用torch.onnx.export底层是通过 tracing追踪或 scripting脚本化来捕获计算图。这个过程有时候会引入一些不必要的操作某些 Python 控制流被展开成冗长的计算图某些 in-place 操作被替换成 copy 操作数据类型转换节点被插入这些都可能导致导出后的模型比原始 PyTorch 模型更慢。导出之后建议用 Netron 之类的工具可视化一下计算图检查有没有异常。五、内存占用模式这是另一个经常被忽略的话题。5.1 模型加载阶段ONNX 模型文件本身就是权重 计算图的序列化。加载时ORT 会反序列化 protobuf解析计算图结构将权重数据加载到内存CPU或显存GPU执行图优化这一步会产生临时的内存开销为中间张量预分配内存池加载完成后的内存占用大致等于模型权重大小 中间张量内存池 引擎自身开销。对比 PyTorchPyTorch 加载模型时只加载权重中间张量是在 forward 时动态分配的。所以 ORT 的初始内存占用通常比 PyTorch 高一些但推理时的内存波动更小。5.2 推理阶段ORT 的内存管理策略和 PyTorch 有本质区别PyTorch动态分配每次 forward 时按需分配中间张量的内存forward 结束后中间张量被 Python GC 回收或者由 CUDA caching allocator 缓存内存占用呈锯齿状波动forward 时上升结束后下降ORT预分配 复用加载时分析整张计算图计算出中间张量的最大内存需求预分配一个内存池推理时所有中间张量从池中分配生命周期不重叠的张量共享同一块内存内存占用基本恒定没有波动5.3 实际影响这种内存模式的差异在以下场景中很重要长时间运行的服务ORT 的恒定内存占用更友好不会因为 GC 延迟导致内存峰值。PyTorch 的锯齿状模式在高并发下可能导致 OOM因为多个请求的内存峰值可能叠加。嵌入式/资源受限环境ORT 的内存占用可预测方便做资源规划。你可以在部署前就精确知道模型需要多少内存。GPU 显存ORT 的 CUDA EP 同样使用预分配策略。如果你在一张卡上跑多个模型ORT 的显存占用更可控。但要注意ORT 默认会预分配较大的显存池可以通过arena_extend_strategy参数调整。5.4 量化对内存的影响如果在 ONNX 上做量化INT8/FP16内存占用会进一步降低精度权重大小相对中间张量大小相对FP321x1xFP160.5x0.5xINT80.25x取决于实现通常 0.25x~0.5x但量化不是免费的午餐精度损失需要评估。对于音频处理这类对精度敏感的场景建议先做 FP16 量化精度损失通常可忽略INT8 则需要仔细验证。六、实践建议最后总结几条实操经验先 profile再决定要不要转 ONNX。用 PyTorch Profiler 看一下推理的时间分布如果 90% 的时间花在 GPU kernel 上转 ONNX 意义不大。如果大量时间花在 CPU 端的调度和数据搬运上转 ONNX 大概率有收益。导出后一定要做数值验证。用相同的输入对比 PyTorch 和 ORT 的输出确保误差在可接受范围内通常 FP32 下 atol1e-5 是合理的阈值。固定 shape 能固定就固定。哪怕要为不同 shape 导出多个模型也比用动态 shape 跑得快。注意 opset 版本。不同 opset 版本支持的算子不同优化效果也不同。一般建议用最新的稳定版本目前是 opset 18-20。ORT 的 Session Options 值得调。graph_optimization_level、intra_op_num_threads、execution_mode这几个参数对性能影响很大不要用默认值就完事了。如果目标是 NVIDIA GPU考虑 TensorRT EP。它会把 ONNX 子图进一步编译为 TensorRT engine通常比纯 CUDA EP 再快 20-50%但首次加载时间会显著增加。参考ONNX 官方规范ONNX Runtime 性能调优文档ONNX Runtime Graph Optimizations
ONNX - 它到底快在哪,又慢在哪
ONNX - 它到底快在哪又慢在哪一、ONNX 是什么ONNXOpen Neural Network Exchange本质上是一种模型的中间表示格式IR类似于编译器里的 LLVM IR。用 PyTorch 训练完一个模型之后模型的计算逻辑是用 Python 描述的跑推理的时候要经过 Python 解释器、PyTorch 的调度器、再到底层的 CUDA kernel。这条链路很长开销不小。ONNX 做的事情是把模型的计算图从 PyTorch 的世界里导出成一个独立的、与框架无关的静态计算图然后交给专门的推理引擎比如 ONNX Runtime去执行。打个比方PyTorch 训练出来的模型像一份 Python 脚本每次执行都要解释器逐行翻译ONNX 导出后的模型像一份编译好的二进制文件直接跑就行。一个典型的 ONNX 文件.onnx里面存的是计算图的拓扑结构哪些算子、怎么连接每个算子的类型和参数Conv、MatMul、Relu 等模型的权重以 protobuf 格式序列化输入输出的 shape 和数据类型这里要注意一点ONNX 本身只是一个格式规范它不负责执行。真正跑推理的是 ONNX Runtime简称 ORT或者其他兼容 ONNX 的推理引擎TensorRT、OpenVINO 等。说用 ONNX 加速实际上是用 ONNX Runtime 加速。二、ONNX Runtime 为什么能快理解了 ONNX 是什么之后关键问题来了为什么换个引擎跑就能快2.1 去掉了 Python 开销PyTorch 推理时即使模型本身的计算是在 GPU 上跑的Python 层面仍然有大量开销每个算子调用都要经过 Python 的函数调度动态图机制意味着每次 forward 都要重新构建计算图GIL全局解释器锁在多线程场景下是个瓶颈ORT 是纯 C 实现的推理引擎模型加载完之后整个推理过程不经过 Python。对于那些计算量不大但算子数量多的模型比如很多小卷积串联的网络光是去掉 Python 开销就能带来可观的提速。2.2 图优化Graph Optimization这是 ORT 最核心的加速手段。静态计算图的好处是引擎可以在推理之前对整张图做全局优化。常见的优化包括算子融合Operator Fusion把多个相邻算子合并成一个。比如 Conv BatchNorm ReLU 三个算子在训练时是分开执行的因为 BatchNorm 在训练和推理时行为不同但在推理时可以融合成一个 kernel。这样做的好处是减少了 kernel launch 的次数也减少了中间结果在显存里的读写。常量折叠Constant Folding如果计算图里有些节点的输入全是常量比如某个 reshape 的 shape 参数那就在加载时直接算好推理时跳过。冗余节点消除去掉那些对输出没有影响的计算节点比如连续的两次 transpose 如果互为逆操作就直接删掉。内存规划Memory Planning分析整张图的数据流提前规划好每个中间张量的内存分配和复用策略避免运行时频繁的 malloc/free。2.3 硬件特化的执行后端ORT 支持多种 Execution ProviderEPEP适用硬件特点CPU EP通用 CPU使用 MLAS 数学库针对 x86/ARM 做了 SIMD 优化CUDA EPNVIDIA GPU调用 cuDNN/cuBLASTensorRT EPNVIDIA GPU进一步将子图编译为 TensorRT engineDirectML EPWindows GPU通过 DirectX 12 调用 GPU支持 AMD/Intel/NVIDIAOpenVINO EPIntel CPU/GPU/VPUIntel 硬件上的深度优化选对 EP 很重要。同一个模型用 CPU EP 和用 CUDA EP 的性能差距可以是数量级的。三、什么时候 ONNX 能加速说了这么多原理落到实际场景里以下情况用 ONNX 通常能获得明显加速3.1 模型结构固定、输入 shape 固定这是 ONNX 最舒服的场景。输入 shape 固定意味着 ORT 可以在加载时就把所有优化做到位包括内存预分配、kernel 选择、算子融合等。典型例子固定长度的音频帧处理比如每次处理 1024 个采样点固定分辨率的图像分类固定序列长度的 NLP 推理3.2 算子数量多、单个算子计算量小前面说了Python 调度开销是按算子数量线性增长的。如果一个模型有几百个小算子比如 MobileNet 这类轻量网络PyTorch 推理时大量时间花在调度上而不是计算上。换成 ORT 之后调度开销几乎为零提速非常明显。实测数据参考一个包含 200 算子的轻量语音增强模型PyTorch 推理耗时 12ms转 ONNX 后 ORT CPU 推理耗时 3ms快了 4 倍。3.3 需要部署到非 Python 环境如果最终要把模型部署到 C 服务、移动端、嵌入式设备上ONNX 几乎是必经之路。ORT 提供 C/C、C#、Java、JavaScript 等语言的 API不依赖 Python 环境。3.4 CPU 推理场景在纯 CPU 推理的场景下ORT 的 MLAS 库针对矩阵运算做了大量 SIMD 优化AVX2、AVX-512、NEON 等通常比 PyTorch 的 CPU 后端基于 MKL 或 OpenBLAS更快尤其是在 batch size 较小的时候。四、什么时候 ONNX 反而更慢这才是很多人踩坑的地方。4.1 动态 shape 场景如果模型的输入 shape 每次都不一样比如变长序列、不同分辨率的图像ORT 的很多优化就失效了无法预分配内存每次推理都要重新规划算子融合的某些模式依赖固定 shape动态 shape 下无法触发某些 EP比如 TensorRT需要为每个新 shape 重新编译 engine第一次遇到新 shape 时会有巨大的延迟PyTorch 的动态图机制天然支持动态 shape反而没有这个问题。实际建议如果必须处理变长输入可以用 padding 把输入对齐到几个固定的 shape 档位比如 256、512、1024然后为每个档位各导出一个 ONNX 模型或者使用 ORT 的动态 shape 支持但接受一定的性能损失。4.2 模型包含不支持的算子ONNX 的算子集opset是有限的。如果模型里用了某些 PyTorch 独有的算子或者自定义算子导出时可能会导出失败被拆解成一堆等价但低效的基础算子组合后者尤其隐蔽。表面上导出成功了但实际推理速度反而变慢了因为原本一个高效的自定义 kernel 被替换成了十几个基础算子的串联。遇到这种情况要么给 ORT 注册自定义算子Custom Operator要么考虑换用其他推理引擎。4.3 模型本身计算量很大瓶颈在 GPU kernel如果模型的瓶颈是大矩阵乘法或大卷积比如 ResNet-152、大型 Transformer那 Python 调度开销相对于 GPU 计算时间来说微不足道。这种情况下PyTorch 和 ORT 调用的底层 kernel 是一样的都是 cuDNN/cuBLAS性能差距很小甚至可能因为 ORT 的 kernel 选择策略不如 PyTorch 的 autotuner 而略慢。简单说模型越大、单次计算越重ONNX 的加速比越小。4.4 导出过程引入了额外开销PyTorch 导出 ONNX 时使用torch.onnx.export底层是通过 tracing追踪或 scripting脚本化来捕获计算图。这个过程有时候会引入一些不必要的操作某些 Python 控制流被展开成冗长的计算图某些 in-place 操作被替换成 copy 操作数据类型转换节点被插入这些都可能导致导出后的模型比原始 PyTorch 模型更慢。导出之后建议用 Netron 之类的工具可视化一下计算图检查有没有异常。五、内存占用模式这是另一个经常被忽略的话题。5.1 模型加载阶段ONNX 模型文件本身就是权重 计算图的序列化。加载时ORT 会反序列化 protobuf解析计算图结构将权重数据加载到内存CPU或显存GPU执行图优化这一步会产生临时的内存开销为中间张量预分配内存池加载完成后的内存占用大致等于模型权重大小 中间张量内存池 引擎自身开销。对比 PyTorchPyTorch 加载模型时只加载权重中间张量是在 forward 时动态分配的。所以 ORT 的初始内存占用通常比 PyTorch 高一些但推理时的内存波动更小。5.2 推理阶段ORT 的内存管理策略和 PyTorch 有本质区别PyTorch动态分配每次 forward 时按需分配中间张量的内存forward 结束后中间张量被 Python GC 回收或者由 CUDA caching allocator 缓存内存占用呈锯齿状波动forward 时上升结束后下降ORT预分配 复用加载时分析整张计算图计算出中间张量的最大内存需求预分配一个内存池推理时所有中间张量从池中分配生命周期不重叠的张量共享同一块内存内存占用基本恒定没有波动5.3 实际影响这种内存模式的差异在以下场景中很重要长时间运行的服务ORT 的恒定内存占用更友好不会因为 GC 延迟导致内存峰值。PyTorch 的锯齿状模式在高并发下可能导致 OOM因为多个请求的内存峰值可能叠加。嵌入式/资源受限环境ORT 的内存占用可预测方便做资源规划。你可以在部署前就精确知道模型需要多少内存。GPU 显存ORT 的 CUDA EP 同样使用预分配策略。如果你在一张卡上跑多个模型ORT 的显存占用更可控。但要注意ORT 默认会预分配较大的显存池可以通过arena_extend_strategy参数调整。5.4 量化对内存的影响如果在 ONNX 上做量化INT8/FP16内存占用会进一步降低精度权重大小相对中间张量大小相对FP321x1xFP160.5x0.5xINT80.25x取决于实现通常 0.25x~0.5x但量化不是免费的午餐精度损失需要评估。对于音频处理这类对精度敏感的场景建议先做 FP16 量化精度损失通常可忽略INT8 则需要仔细验证。六、实践建议最后总结几条实操经验先 profile再决定要不要转 ONNX。用 PyTorch Profiler 看一下推理的时间分布如果 90% 的时间花在 GPU kernel 上转 ONNX 意义不大。如果大量时间花在 CPU 端的调度和数据搬运上转 ONNX 大概率有收益。导出后一定要做数值验证。用相同的输入对比 PyTorch 和 ORT 的输出确保误差在可接受范围内通常 FP32 下 atol1e-5 是合理的阈值。固定 shape 能固定就固定。哪怕要为不同 shape 导出多个模型也比用动态 shape 跑得快。注意 opset 版本。不同 opset 版本支持的算子不同优化效果也不同。一般建议用最新的稳定版本目前是 opset 18-20。ORT 的 Session Options 值得调。graph_optimization_level、intra_op_num_threads、execution_mode这几个参数对性能影响很大不要用默认值就完事了。如果目标是 NVIDIA GPU考虑 TensorRT EP。它会把 ONNX 子图进一步编译为 TensorRT engine通常比纯 CUDA EP 再快 20-50%但首次加载时间会显著增加。参考ONNX 官方规范ONNX Runtime 性能调优文档ONNX Runtime Graph Optimizations