ARTICLE DETAIL

资讯详情

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

深度学习模型推理加速实战:算子融合与INT8量化从原理到落地

深度学习模型推理加速实战:算子融合与INT8量化从原理到落地 几个月前团队里一个算法同学跑过来找我说他的目标检测模型在CPU上单张推理要六十多毫秒生产环境的GPU资源又紧张问我能不能在不换框架的前提下把速度提上来。我当时手里刚好在折腾一个叫“Model-Optimizer”的小工具本来是自己内部实验用的拿过来一跑单张推理直接降到二十毫秒出头。后来我花了几个周末把它整理成独立项目今天把里面的核心思路、设计取舍和实操过程完整写出来希望对正在做模型部署和推理加速的同行有点帮助。Model-Optimizer本质上是一个面向深度学习模型推理阶段的轻量级优化工具链。它接收ONNX格式的模型作为输入通过算子融合、INT8量化、图结构优化和内存规划四类手段输出一个推理速度更快、体积更小、内存占用更低的优化后模型。适合三类人参考一是做模型部署的工程师二是想搞懂推理加速原理的算法工程师三是刚入门、想找一个可落地的模型优化项目的开发者。1. 为什么做Model-Optimizer模型能跑但不代表能上线1.1 项目起点一个被推理延迟卡住的上线需求很多团队都有这个问题模型在训练机上精度刷得漂漂亮亮一进生产环境就露馅。训练时PyTorch自带的高效算子掩盖了模型结构里的冗余多Batch训练也掩盖了单条推理时的显存和延迟问题。真正上线时模型面临的是一次只推理一张图、CPU负载有限、内存带宽受限的残酷环境。当时那个检测模型就是这个情况。模型本身是个YOLOv5s结构训练时跑得挺好但到了CPU端的部署环境单张图片推理耗时六十多毫秒一秒只能处理十几张。业务方要求至少每秒处理二十张这就意味着必须在不改变模型效果的前提下把推理延迟砍掉一大截。我先试了常规手段ONNX Runtime的图优化开满、线程数调高、输入分辨率保持最小合法值。结果最好也就到四十多毫秒离目标还差一半。这让我意识到通用框架的优化是“黑盒”的它不会替你精细处理每一个算子的计算模式更不会针对特定硬件做激进的融合和量化。Model-Optimizer的雏形就是从“我要能自己控制每一步优化”这个念头开始的。1.2 Model-Optimizer到底是什么简单来说Model-Optimizer是一个把ONNX模型“重新编译”一遍的工具。它包含三层第一层是前端负责读取ONNX模型做基本的合法性检查和结构解析。第二层是优化引擎内部按固定顺序跑一系列优化Pass每个Pass只做一件事比如“把卷积后面的批归一化折叠进去”或者“把连续两个重算子合并成一个”。第三层是量化模块和运行时量化模块负责把FP32模型转成INT8运行时则负责在目标硬件上高效执行优化后的模型。这个项目我刻意没有做成“全自动黑盒优化”因为全自动意味着不可控不可控意味着上线出问题难排查。Model-Optimizer反而保留了中间结果的导出能力每一步优化后模型都能单独导出并对比精度和速度。哪个Pass出了问题直接把对比结果拉出来看就行。项目的核心价值可以概括成三句话推理变快、模型变小、内存占用更可控。实测效果后面会展开我先说数据一个ResNet50模型FP32下CPU推理一百二十毫秒优化后四十五毫秒模型体积从二十五兆降到六点五兆内存峰值从优化前的一百二十兆降到三十五兆。量化精度损失控制在百分之一以内Top-1准确率下降约0.8%。1.3 为什么没有直接选TensorRT或ONNX Runtime的现成方案很多朋友会问TensorRT不是现成的吗ONNX Runtime也有GraphOptimizationLevel直接用不就行了这个问题我认真权衡过。TensorRT确实是NVIDIA GPU上的最优解之一但它的限制也很明显一是绑定NVIDIA硬件CPU、AMD GPU、ARM设备都用不了二是TensorRT对模型算子有严格约束模型里出现不支持的算子就得拆图或者手写插件维护成本很高三是TensorRT的优化过程更接近“黑盒”出了性能问题很难精细定位是哪一层引起的。ONNX Runtime的图优化虽然覆盖场景广但对优化细节的控制力不够。比如它默认的融合策略里ConvBNReLU会融合但你没法单独控制某个Pass的行为也没法查看融合后每一层的中间输出做精度对比。一旦量化精度掉了你只能整体回退缺少精细定位的能力。Model-Optimizer走的是一条更“可控”的路线不追求在所有硬件上做到极限性能而是保证在常见CPU和GPU上拿到稳定收益同时让使用者能看清楚每一步优化做了什么、代价是什么。对于中小团队来说这种“透明”比“极限”更重要。2. 核心优化手段拆解Model-Optimizer到底做了什么2.1 算子融合把“跑不满”的算子捆在一起算子融合是所有优化手段里见效最快的一项也是Model-Optimizer最核心的优化Pass。它的原理用一句话说就是把多个连续的小算子合并成一个大算子减少中间结果的读写内存次数。先说ConvBN融合这个在推理阶段是数学上完全可折叠的。BN在推理时是一个逐通道线性变换公式是y (x - mean) / sqrt(var eps) * gamma beta而卷积是线性运算所以BN可以完全折叠进卷积的权重和偏置里。折叠后权重变为 w w * gamma / sqrt(var eps)偏置变为 b (b - mean) * gamma / sqrt(var eps) beta。这样一来原来需要两次算子调用、两次内存读写现在只需要一次计算量虽然没变但内存访问的开销省了一大截。如果再往下看卷积后面往往还跟着一个ReLU激活ReLU的处理逻辑也能一并吸收到卷积的output写回阶段这就是ConvBNReLU三段融合。实测中这个融合能让单层耗时平均下降百分之二十到三十。Model-Optimizer还实现了ConvBNSiLU、AddReLU、Gelu近似融合等模式。Gelu在CPU上是一个偏贵的操作如果模型里大量使用我会用ReLU近似或者一个多项式逼近替换掉精度影响很小但速度提升明显。这也是一个值得单独说的点模型优化不只是数学等价变换在精度损失可接受的前提下做一些“工程近似”是完全值得的。2.2 INT8量化压缩的核心手段模型体积和推理耗时的最大头其实来自数值精度太高。FP32的模型在推理时大量的时间浪费在无关紧要的精度上。INT8量化就是把模型的权重和激活值从32位浮点压到8位整数计算量、内存带宽和模型体积都会成倍下降。Model-Optimizer采用的是训练后量化PTQ方案流程分四步第一步用一小批真实数据跑一遍FP32模型统计每一层激活值的分布第二步根据分布确定每一层的量化参数scale和zero_point第三步把权重和激活值映射到INT8范围第四步做精度对比必要时对敏感层回退为FP32。量化参数的计算是核心。以最常用的对称量化为例假设某一层激活值经过统计范围在[-0.6, 3.2]那么先取绝对值最大值3.2然后计算scale 3.2 / 127 0.0252。一个浮点值3.0对应整数1193.0 / 0.0252四舍五入。推理时整数119乘回scale就能还原近似值。这个计算看起来简单真正容易踩坑的地方在于激活值分布里可能有个别极端值会把scale拉大导致普通数值损失更多精度。所以Model-Optimizer的校准策略默认使用Percentile百分位统计而不是MinMax。比如取百分之九十九点九九的分布范围把最极端的少量点截断掉这能显著改善量化后精度尤其是检测和分割模型。校准数据不需要多两百张左右能代表真实业务分布的图就够了后面我会详细说。在量化粒度上卷积权重默认用per-channel量化也就是每个输出的通道各算一组scale原因是不同通道的权重分布差异可能很大共用一组scale会导致某个通道严重失真。激活值因为卷积输出是跨通道计算的用per-tensor更稳妥也不太会增加运行时复杂度。2.3 图优化与常量折叠除了融合和量化还有一批“机械性”优化手段。这类优化藏在模型结构中最不起眼的地方但累计收益不可小视。常量折叠是最容易理解的一类优化模型里有些算子的输出其实在推理过程中是固定的比如Const节点的Add、Mul、Padding、Slice、Gather等操作它们的输入都是常量输出完全可以在优化阶段算好直接替换成Const节点。这样推理时就能省掉一个算子的执行时间。这个看起来很简单但很多从PyTorch导出的模型里这类冗余还真不少尤其是Padding和Reshape附近。另一个高频优化是冗余结构消除。PyTorch导出ONNX时受限于Trace机制经常会出现一连串的Reshape、Transpose、Squeeze、Unsqueeze。有些是必要的维度转换有些则是纯冗余比如两个连续的Transpose又转回原样或者某个Reshape前后的shape根本没变。Model-Optimizer会逐个检查这些维度变换节点能删的直接删能合并的合并。公共子表达式消除也值得一提。如果同一个大Tensor在模型中被多个分支使用而这些分支里恰好都包含相同的一个子计算图那么这个子图只需要计算一次结论共享即可。这个优化在自注意力结构里尤其有效比如Q、K、V向量在主分支中可能会被重复做若干次同样的投影操作实际上只需要算一次后切片。图优化看起来不起眼但这些叠加起来通常能带来百分之十到十五的额外加速。更重要的是做完图优化之后再量化量化统计的分布会更稳定因为残留的冗余算子都被清掉了。2.4 内存规划与复用内存优化往往是被忽视的一环但它在实际部署中非常关键。训练框架习惯了大显存可到了端侧内存是论兆算的。Model-Optimizer在优化阶段会对模型做一次完整的数据流分析记录每个Tensor的生命周期也就是从哪个算子开始被使用到哪个算子之后不再被使用。然后根据生命周期做缓冲区复用规划。如果两个Tensor的生命周期不重叠它们可以共享同一块内存空间。这个思路和操作系统里的内存分配是类似道理。我用一个PagedArena的内存分配器来管理推理时的内存它按页预分配大块内存然后依据Tensor生命周期逐步划分。实践下来ResNet50这种链式结构模型激活值的理论总和在一百二十兆左右但实际推理时的峰值内存只用三十五兆就够了节约了约七成。Batch一大的时候收益更明显。一个Batch为8的检测模型原始峰值内存接近六百兆优化后压到一百八十兆以内。这个优化在CPU和端侧NPU上尤其有用直接决定了模型能否跑得起来。3. 实操全过程从PyTorch模型到加速后的部署模型3.1 环境准备与安装Model-Optimizer的运行环境要求不高Python 3.8以上安装了ONNX和NumPy的基本环境就能跑。核心优化引擎我为了方便跨平台用Python实现的Pass框架加上C扩展的量化内核后者只在执行INT8推理时才需要编译。安装方式就一条命令pip install model-optimizer如果是从源码构建需要额外安装CMake和C编译器。建议直接用pip安装省事。在开始之前还需要准备一个ONNX模型。如果你手里只有PyTorch的.pt权重先导出成ONNX再继续操作这一步后面会专门说。3.2 输入模型准备导出ONNX与基础检查我强烈建议在开始优化前先把ONNX模型本身的“健康状况”搞清楚。用ONNX Runtime跑一遍确认输入输出正常、精度符合预期再交给Model-Optimizer。从PyTorch导出时我习惯用固定opset版本一般选13到17之间的某个稳定版本。太老的opset会导致模型里残留一些不规范的算子表示增加优化难度。导出命令参考import torch import torch.onnx model torch.load(yolov5s.pth, map_locationcpu) model.eval() dummy_input torch.randn(1, 3, 640, 640) torch.onnx.export( model, dummy_input, yolov5s.onnx, opset_version13, input_names[images], output_names[outputs], dynamic_axes{images: {0: batch_size}, outputs: {0: batch_size}} )导出后第一时间用onnx.checker.check_model和onnx.shape_inference.infer_shapes做一次基础检查这两个是ONNX官方工具能发现问题模型的大多数结构性问题。3.3 一键优化算子融合与图结构精简Model-Optimizer的入口是一个命令行工具。最基础的使用方式model-optimizer optimize yolov5s.onnx -o yolov5s_opt.onnx执行时会顺序跑下面这些Pass[0/6] Shape inference [1/6] Constant folding [2/6] Redundant transpose elimination [3/6] Conv-BN fusion [4/6] Conv-ReLU fusion [5/6] Buffer lifetime analysis每一步执行完都会打印统计信息。比如Conv-BN融合后模型里的算子数量从四百二十个降到三百二十五个减少约百分之二十。算子数量降下来了推理引擎的任务调度开销就小很多。这里的优化顺序是有讲究的先做常量折叠和冗余消除再做融合。因为如果模型里残留大量冗余Transpose它们会干扰Conv的输入形态判断可能导致融合误判。这就是我前面说的“先图优化再量化”每一步优化都依赖前一步把模型结构理得更干净。优化完成后建议用onnxruntime或model-optimizer evaluate验证一下优化前后模型的输出是否一致。对非量化的纯图优化理论上输入输出应该完全一致最多在浮点误差范围内有细微差异。如果差异太大通常是某个Pass实现有问题需要用二分法逐层排查。3.4 INT8量化与校准图优化完成后就可以进行INT8量化了。Model-Optimizer会把量化分成两个阶段。第一阶段是校准。需要准备一小批校准数据推荐从真实业务数据中采样我这里用的是验证集的子集约两百张图。校准的目的是统计模型每一层激活值的真实分布而不是拍脑袋给一个范围。model-optimizer quantize yolov5s_opt.onnx \ --calib-dir ./calib_imgs \ --calibrate-samples 200 \ --quantize-mode int8 \ -o yolov5s_int8.onnx在量化期间程序会逐层输出当前这层的激活值范围。有经验的工程师会特别留意那些分布明显不正常的层比如有些层的激活值几乎全集中在零附近只有偶尔几个大值。这类层往往是最敏感的建议配合“敏感层跳过”功能把这部分层保持FP32其他层走INT8。Model-Optimizer提供了--skip-layer参数可以指定按名字跳过某层。第二阶段是产出版本。量化结束后模型体积会有直观变化YOLOv5s原始ONNX是14.8兆INT8后变成4.1兆降幅约百分之七十二。模型体积下降对网络传输和磁盘占用的改善是立竿见影的。3.5 精度与速度验证量化完成后最关键的验证环节是精度和速度的对比。Model-Optimizer内置了一份评估工具支持对比任意两个模型的逐层输出、指定指标精度和推理时延。先单张图对比逐层输出的余弦相似度model-optimizer evaluate-diff yolov5s.onnx yolov5s_int8.onnx \ --input-img ./test.jpg \ --metric cosine输出结果里会列出每一层两个输出版本的余弦相似度。大多数层应该在0.99以上个别层如果低于0.95就要重点关注。然后是整体指标精度验证我一般用一段独立的Python脚本检测mAP对比FP32和INT8版本的指标差异。最终结果如下指标FP32原始模型优化后INT8模型变化平均推理延迟64.2 ms21.4 ms下降66.7%模型体积14.8 MB4.1 MB下降72.3%mAP500.5410.532下降0.9%峰值内存约420 MB约170 MB下降59.5%这个结果不是我最优的一次但属于很典型的收益曲线。INT8量化几乎不影响精度却换来接近三倍的推理速度提升。速度测量时要注意先把前几次推理作为预热丢弃从第十次之后开始计时取平均值否则缓存未命中会让结果虚高。4. 常见问题与排查技巧实录4.1 量化后精度掉点严重怎么办这是INT8量化最让人头疼的问题。如果整体精度跌了超过百分之二先不要慌大概率不是全局问题而是少数几个“敏感层”带崩了全局。排查方法是做逐层敏感度分析把模型每一层分别做量化其他层保持FP32逐层对比输出和基线之间的误差把误差最大的几层标记出来。Model-Optimizer提供了半自动的敏感层搜索功能model-optimizer analyze-sensitivity yolov5s_opt.onnx \ --calib-dir ./calib_imgs \ --topk-sensitive 5找到敏感层后两种应对方案一是把这些层从量化列表中移除保持FP32二是如果敏感层是卷积尝试从per-tensor量化为per-channel量化。实测中百分之八十的精度掉点问题靠这两招都能解决。还有一个容易被忽视的原因是校准集质量。如果校准集和真实业务数据分布差异很大统计到的激活值范围就不准量化损失自然大。我遇到过用ImageNet校准集给工业质检模型做量化精度掉百分之四换成业务现场的五千张真实图片后掉点缩小到百分之零点五以内。4.2 算子不支持导致转换中断导出ONNX时PyTorch支持的算子非常多但有些算子在不同opset版本里表示方式差异极大。最常见的报错是某个算子版本不兼容或者新出现的算子没有对应实现。如果遇到算子不支持的报错按下面的顺序排查第一升级ONNX opset版本重新导出模型。PyTorch新版支持更高的opset很多老版本里手工拼接的结构在新版本中有了规范表示。第二用ONNX Simplifier做一次模型化简它会自动替换掉一批PyTorch导出的冗余结构。第三如果模型里有太冷门的自研算子考虑在PyTorch导出前改代码用手工等价结构替代。Model-Optimizer在遇到不支持的算子时也会把完整上下文打印出来包含算子的类型、输入输出shape、来自原始模型哪个位置。这时候不要盲目删节点先看结构大部分自定义算子在ONNX标准op集里都能找到替代品。4.3 动态shape问题用dynamic_axes导出的模型在运行时输入可能变化。Model-Optimizer对动态shape的调用方式有点特殊优化阶段需要给定一个参考shape用于静态内存规划但运行时仍能接受一定范围内的动态shape。model-optimizer optimize yolov5s.onnx \ --batch 1 \ --height 640 \ --width 640 \ -o yolov5s_opt.onnx这里的参数会用来做内存规划与算子融合验证不会把模型结构钉死成静态shape。真正实现时需要给动态维度的范围设置一个上限超过上限会触发重规划。一个容易踩的坑是动态shape模型的优化结果在某个shape下OK换个shape精度没问题但速度变慢。这是因为内存复用方案是依赖具体shape的在优化时锁定了shape范围。解决办法就是给动态维度的上限留一些余量避免运行时频繁触发重新规划。4.4 CPU上INT8提速不明显有朋友反馈在CPU上跑INT8量化后的模型速度几乎没有提升。这个问题涉及两个层面的原因。第一个原因是推理内核没有走到支持INT8的路径。CPU推理性能依赖SIMD指令集比如AVX2、AVX512。如果推理没有用对编译选项或者算子实现里没有针对INT8数据类型的特化kernel实际计算时可能还是先把INT8转回FP32再算那就白优化了。Model-Optimizer内置的运行时默认开启了AVX2支持但我建议在真实部署前做一次指令集检测确认目标机器的CPU支持情况。第二个原因是模型太小、线程太多。一个只有几十个算子的轻量模型开启八线程时线程调度和同步的开销可能比计算本身还大。这个问题的信号是线程数从一调到八推理时延不降反升。解决办法是使用自动线程自适应根据模型运算量和CPU核数自动选择最优线程数而不是无脑调大。4.5 常见问题速查表现象可能原因解决办法量化后精度掉点多有敏感层用敏感度分析定位跳过敏感层或改为per-channel校准后模型精度漂移校准集分布不符换成真实业务数据校准集200张以上某个节点报错不支持ONNX opset版本旧、算子表达不规范升级opset、用onnx-simplifier、手工替换为等价op动态shape推理变慢内存规划被反复触发限制shape variation范围预留缓存CPU上INT8无提升SIMD未启用、kernel未特化、线程配置不当检查CPU指令集、开启AVX2、自动自适应线程数优化后输出不一致某个Pass实现有bug逐步关闭Pass用二分法定位出错Pass内存峰值反而变大缓冲区复用策略失效检查是否还有动态shape在修改运行时规划5. 写在最后的几点实操心得在这些优化项目里摸爬滚打下来最深的体会是“优化不是堆技巧而是建立可验证的流程”。每次做一个Pass之前我先想清楚怎么验证它再动代码。量化前先跑一遍预测脚本记下每层输出哈希融合后立刻对比哈希是否一致。有了这个基线出问题就是在某个Pass内部不用来回猜。另外一个小技巧校准集不用贪多但一定要覆盖真实场景中的极限情况。如果是白天黑夜都要用的检测模型校准图里就两种光照都得有如果有运动模糊就放几张运动模糊的图。校准集的价值在于反映真实的数据分布不在于数量堆得多大。这个项目后续我计划扩展的方向有三个一是稀疏化与半结构化剪枝配合INT8量化能进一步压缩模型二是低比特量化INT4、FP16混合在部分场景下收益会比INT8更明显三是自动化的Pass组合搜索根据目标硬件特性自动选择最优的优化组合。如果你也在做模型部署相关工作欢迎把Model-Optimizer拉下来跑一下上面的坑和方法都是我真实踩出来的能让你少走不少弯路。
返回列表