资讯详情

资讯详情

Model-Optimizer:量化、剪枝、蒸馏到ONNX的模型优化流水线实践

最近一直在折腾模型交付的事训练好的模型推上线体积大、延迟高领导一句话再优化优化就得从早到晚对着各种工具来回试。说实话模型优化这件事单点方案网上到处都是但真正能把量化、剪枝、蒸馏、导出这些步骤串成一条顺手的流水线每一步都把坑踩平才是最有价值的部分。这套被我命名为 Model-Optimizer 的实践流程就是干这个的这篇文章把完整思路和关键细节都摊开讲适合正在做模型上线、推理加速或者部署方案选型的同学参考。如果刚接触部署也能照着这套流程把优化这件事完整跑通。Model-Optimizer 解决的核心问题说白了就是让一个已经训练好的模型以更小的体积、更快的速度和更低的资源占用运行起来同时尽量保住原有精度。它的主线是量化、剪枝、知识蒸馏、结构重参数化然后把优化后的模型统一导出成 ONNX交给推理引擎做最后一步加速。这四件事每一件单拿出来都有很多资料但真正把它们串成一条完整流程、把每一步的边界条件都处理好才是实际工程里最花精力的地方。1. Model-Optimizer 整体设计为什么把优化流程做成一条流水线刚开始接手模型优化时面对的情况和大多数团队差不多模型文件五花八门有 PyTorch 训出来的也有 TensorFlow 导出的部署环境有 CPU 服务器、有 GPU 卡、还有边缘设备业务对时延要求也各不相同有的要求单次推理 50 毫秒以内有的能容忍 1 秒。如果每次优化都临时找一个方案拼凑出来的流程既不可复现出了问题也难排查。1.1 从一次性优化到可复用流水线的思路转变我得到的第一个教训是不要一上来就追求某个具体优化手段比如我要用 INT8 量化而是先想清楚全链路长什么样。Model-Optimizer 的流水线设计在第一版就确定了几个原则输入统一不管上游框架是什么先转成统一的中间表示再进入优化环节。步骤解耦量化、剪枝、蒸馏彼此独立可以单独启用或关闭方便做消融实验对比每步收益。出口统一所有优化后的模型必须能导出到 ONNX因为 ONNX 是当前推理引擎兼容性最好的中间格式。指标前置每一步优化前先固定评测脚本和数据集确保每一步的精度变化是可量化的。这套设计有个很大的好处当新模型进来时优化流程是确定的而不是靠灵感。团队里任何人接手都能照着流水线跑一遍产出的报告格式也一致。1.2 为什么去掉多引擎适配只围绕 ONNX 做出口早先版本里我想让流水线同时输出 TorchScript 和 ONNX 格式结果维护成本极高。TensorRT 的优化需要 ONNX 图中有符合要求的算子OpenVINO 也需要单独的模型格式转换如果每一步都去适配工作量会翻倍。最终我把出口统一到 ONNX推理侧交给不同引擎各自处理。这里用一个生活化的类比解释为什么要这么做ONNX 之于模型部署就像通用插座之于家电。电器上只要有一个统一的三孔插座无论发电厂在哪儿你都能把电接过来用。模型在这个统一格式上可以被不同的推理引擎各自优化而不需要每种引擎都去重写一套转换逻辑。提示如果你的模型是纯 PyTorch 且部署环境固定也可以考虑直接跳过 ONNX用 TorchScript 或模型原生方式部署。这个取舍适合模型只有一个、部署栈完全固定的场景一旦模型增多、环境复杂还是走 ONNX 更省事。2. 核心优化手段详解量化、剪枝、蒸馏的作用与选型很多人对模型优化的理解就是量化一下模型变小了速度变快了。这句话没错但真做起来会踩到不少坑。这里详细拆解一下几个核心手段的原理和适用边界。2.1 量化把 FP32 变成 INT8精度和速度的平衡术量化是最直接见效的一步。它的核心思想很简单训练和推理时模型权重通常用 FP32 浮点数存储每个数占 4 字节如果换成 INT8 整数每个数只占 1 字节模型体积直接缩小到四分之一。更重要的是INT8 计算在 CPU 上可以利用很多硬件加速指令推理速度往往比 FP32 快 2 到 4 倍。但是量化付出的代价是精度损失尤其是大模型和小模型损失程度不一样。怎么理解这一步呢就好像你手里有一张原始分辨率极高的照片用 JPEG 压缩之后肉眼看上去差别不大但文件体积大幅减小如果压缩过度就会出现可见的模糊色块。模型量化后的压缩痕迹就表现为精度下降有时候是小数点后几位有时候直接掉几个点。量化有两种主流方式实际项目中要区分开训练后动态量化只对权重做 INT8 量化激活值在推理时动态计算范围。实现简单CPU 上很有效适合 RNN、Transformer 这类模型通常不需要校准数据。训练后静态量化权重和激活值都做 INT8 量化需要准备一小批校准数据来统计激活值的分布效果通常比动态量化更好但需要额外的校准流程。实际项目中我做过一个 BERT 分类模型的动态量化模型体积从 420MB 降到约 110MB推理时延从 80 毫秒降到 40 毫秒左右精度损失不到 0.3 个点。这个收益相当可观说明动态量化是性价比最高的起步方案。2.2 剪枝去掉模型里的赘肉另一个思路是剪枝。深度学习模型里往往存在大量冗余参数有些神经元或权重对最终输出贡献很小。剪枝的思想是把这些低贡献的部分直接去掉留下一个更稀疏、更小的网络。还是用生活类比剪枝就像修剪一棵树不是把树砍小而是去掉那些不影响主干生长的枝丫。模型经过剪枝后参数量显著减少对显存和算力的需求也跟着下降。剪枝的粒度可以很粗也可以很细非结构化剪枝直接在权重矩阵上把接近零的值置零得到稀疏矩阵。这种剪枝对存储有利但需要专用的稀疏库才能加速对通用推理引擎不太友好。结构化剪枝按行、列或通道粒度删除权重比如直接剪掉某些卷积通道。这样得到的是规整的稠密矩阵可以被通用推理引擎直接加速实践中应用更广。实际使用时结构化剪枝通常需要微调恢复精度。剪枝比例要控制好一般我先从 10% 开始测试逐步递增找到精度下降可接受范围内的最大剪枝率。2.3 知识蒸馏让大模型当老师教出一个小模型如果说量化和剪枝是在压缩现有模型知识蒸馏则是在重新训练一个更小的模型。核心思路是训练一个小模型模仿大模型的输出。大模型老师的输出不仅有标签信息还有对各类别的置信度分布这些软标签包含大量暗知识小模型学到的内容会比直接训练更丰富。实际应用时一个常见组合是用蒸馏把模型从轻量级 BERT 蒸馏到更小的结构再用量化把蒸馏后的小模型进一步压缩。这种组合在 NLP 任务上尤其常见比如把 BERT-base 蒸馏成 TinyBERT参数量可以缩小到原来的几分之一精度还能保持得很高。蒸馏训练的代价是需要额外的训练时间和数据项目时间紧时我通常把蒸馏放在优化步骤的后端先看量化和剪枝能不能满足需求不能满足再考虑蒸馏。2.4 结构重参数化推理阶段折叠卷积层还有一个容易被忽略但很实用的技巧是结构重参数化在 PyTorch 里表现为重参数化或折叠 BN。训练过程中模型往往带有 BatchNorm 层便于训练稳定但推理阶段 BN 层可以折叠进卷积层减少计算量。把 BN 层折叠进前一层卷积不仅省去一次张量运算还能减少几次 GPU kernel launch 的开销。这个优化很轻量几乎不损失精度刚开始时我也经常忽略它后来发现它对小模型的加速效果还挺显著。Model-Optimizer 的流水线里这一项默认开启。3. 实操在 Model-Optimizer 流水线上跑通一个 BERT 分类模型前面讲了原理和选型这一节进入真正的实操。以一个实际跑过的 BERT 中文分类模型为例把整个流程的每个环节都记录下来。这个模型训练时用的是 Hugging Face 的 transformers 库任务是对客服对话做多标签分类。部署目标是单核 CPU 服务器要求模型内存占用小于 500MB单次推理时延小于 100 毫秒。背景参数先列一下项目数值原始模型结构bert-base-chinese12 层 Transformer约 1.02 亿参数原始模型体积约 420MBFP32 权重评测数据集5000 条客服对话样本精度指标宏平均 F1部署环境4 核 CPU16GB 内存无 GPU3.1 第一步建立基线别在优化前把评价标准弄丢了正式开始优化之前必须先跑原始模型的基线指标。这一步看起来简单但很多人图快直接跳过结果优化完发现自己根本不知道相对于什么水平在提升。用 transformers 加载模型写一个简单的评测脚本输出 F1 值和平均推理时延。import time import torch from transformers import AutoTokenizer, AutoModelForSequenceClassification from sklearn.metrics import f1_score model_name bert-base-chinese tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForSequenceClassification.from_pretrained(model_name) model.eval() def evaluate(model, dataloader): model.eval() preds, labels [], [] start time.time() with torch.no_grad(): for batch in dataloader: outputs model(**batch) preds.extend(torch.argmax(outputs.logits, dim-1).tolist()) labels.extend(batch[labels].tolist()) elapsed time.time() - start return f1_score(labels, preds, averagemacro), elapsed / len(dataloader) f1, latency evaluate(model, val_loader) print(f基线模型 F1{f1:.4f}, 平均时延{latency:.4f}s)基线结果 F1 约 0.9143单条推理平均约 78 毫秒。这个数值就是后面每一步优化参照的锚点。3.2 第二步执行 Post-Training Dynamic Quantization在 PyTorch 里一行代码就能完成 BERT 的动态量化。quantized_model torch.quantization.quantize_dynamic( model, {torch.nn.Linear}, dtypetorch.qint8 )这里有个关键点为什么只量化 Linear 层因为 BERT 主要由 Embedding、LayerNorm 和 Linear 组成动态量化目前主要支持 Linear、LSTM 等层Embedding 和 LayerNorm 会被保留为 FP32。这样量化的覆盖范围虽然不完整但 Linear 层占了 BERT 绝大部分参数收益已经很可观。量化完成后再次运行评测脚本指标基线模型动态量化模型模型体积420MB约 110MB平均时延78ms41ms宏平均 F10.91430.9127体积少了近四倍时延快了近一半F1 只掉了 0.0016。这次量化几乎没什么代价。这验证了我之前的习惯优先做动态量化性价比太高了。3.3 第三步导出 ONNX用推理引擎做二次优化量化完成后的模型还是 PyTorch 格式要部署到 CPU 服务器上我选择导出成 ONNX再用 onnxruntime 进行推理加速。导出 ONNX 的代码import torch from transformers import AutoTokenizer sample_input tokenizer(你好这是一句测试文本, max_length128, paddingmax_length, truncationTrue, return_tensorspt) onnx_path bert_quantized.onnx torch.onnx.export( quantized_model, (sample_input[input_ids], sample_input[attention_mask]), onnx_path, opset_version14, input_names[input_ids, attention_mask], output_names[logits], dynamic_axes{ input_ids: {0: batch_size, 1: seq_len}, attention_mask: {0: batch_size, 1: seq_len}, logits: {0: batch_size} } )有几个细节必须提醒dynamic_axes很重要不设置的话导出的模型只支持固定长度输入实际部署时一旦传入更长的文本就会报错。可以在导出时就把 batch_size 和 seq_len 设成动态维度代价是推理时稍微多一点点开销但换来使用上的灵活。opset_version建议选 14 以上太老的 opset 对 Transformer 里的一些算子支持不完整导出过程容易失败。导出的模型文件虽然已经量化过但 ONNX Runtime 的 CPU 执行器对量化模型的支持取决于它使用的优化能力。可以不额外调用量化 API直接让 onnxruntime 运行它也能感知模型内已有的量化结构。3.4 第四步onnxruntime 推理实测线程数其实是个大坑导出 ONNX 后用 onnxruntime 加载并测试性能import onnxruntime as ort import numpy as np import time sess_options ort.SessionOptions() sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_options.intra_op_num_threads 4 session ort.InferenceSession(onnx_path, sess_optionssess_options, providers[CPUExecutionProvider]) def ort_infer(input_ids, attention_mask): outputs session.run( [logits], {input_ids: input_ids, attention_mask: attention_mask} ) return outputs[0] input_ids np.array([sample_input[input_ids].numpy()[0]]) attention_mask np.array([sample_input[attention_mask].numpy()[0]]) # 预热推理排除第一次加载时初始化开销 for _ in range(5): ort_infer(input_ids, attention_mask) times [] for _ in range(100): t0 time.time() ort_infer(input_ids, attention_mask) times.append(time.time() - t0) print(f平均推理时延: {np.mean(times)*1000:.2f} ms)实测结果出来后onnxruntime 在后端做了不少算子融合优化时延从 PyTorch 动态量化的 41ms 进一步下降到了 28ms 左右比最初的 78ms 快了接近三倍。这里要特别提醒线程数的影响一开始设成 8 线程结果时延反而涨到了 35ms。原因在于该模型的单次计算量并不算特别大8 线程的调度开销超过了并行计算带来的收益。这个例子说明不是线程越多越快建议实际压测时把线程数从 1 到 8 逐个跑一遍才能找到当前模型在目标机器上的最优配置。3.5 第五步验证精度与延迟的最终状态优化流程末尾用相同评测集重新跑了完整评测指标基线模型动态量化模型 ONNX Runtime模型体积420MB约 110MB平均时延78ms28ms宏平均 F10.91430.9130精度损失约为 0.0013时延优化超过 60%体积压缩到 26%。当前方案完全满足上线约束。这一套流程下来最耗时的是导出 ONNX 和排查算子兼容性真正量化推理反而很快。4. 进阶优化静态量化与算子兼容性处理如果动态量化后仍然无法达到部署要求例如时延目标更严格或者模型部署在资源非常有限的边缘设备上就需要升级到静态量化。静态量化比动态量化更复杂收益也更大但它也常常和各种算子兼容性问题交手。4.1 静态量化的完整流程与校准数据准备静态量化除了量化权重外还需要量化激活值。因为激活值在前向推理中是被算出来的数值范围需要提前统计所以需要一小批有代表性的校准数据。# 伪代码静态量化的大致流程 model.qconfig torch.quantization.get_default_qconfig(fbgemm) torch.quantization.prepare(model, inplaceFalse) # 遍历校准数据集喂入模型收集激活值分布统计信息 for batch in calibration_loader: model(batch) torch.quantization.convert(model, inplaceFalse)校准数据通常是几百到几千条即可但要有代表性必须覆盖业务真实场景中的输入分布比如文本分类任务中各个类别都要有样本否则统计出来的激活值范围不准确推理时遇到没见过的范围就会发生严重精度损失。我做多标签分类校准的时候选的是每个类别至少 50 条样本整批校准集 2000 条大概花了 5 分钟。静态量化后同一个模型在 CPU 上的运行时间进一步降到 21ms 左右模型体积又缩到约 100MB但 F1 从 0.9130 掉到了 0.9108。相比动态量化静态量化精度损失确实更明显优势是速度更快、部署端更稳定。4.2 算子不兼容问题清单与处理方法在导出 ONNX 和 onnxruntime 推理时算子兼容性是新建工程最常见的烦恼。这里把实际遇到过的问题列成一张速查表问题现象可能原因处理方法导出时报Unsupported operator模型的某些自研算子没有对应的 ONNX 实现为算子注册自定义符号或用周边等价算子组合替代onnxruntime 加载时报内存错误动态轴设置错误或输入维度不匹配核验 input shape 与模型签名固定维度测试推理结果全为 NaN量化时校准数据分布覆盖不足增加校准集规模并检查是否包含边缘样本推理结果正确但速度很慢线程数配置不当或图优化级别偏低尝试不同线程数把 graph 优化级别设为 ORT_ENABLE_ALL导入到 TensorRT 时算子不支持ONNX 中的部分算子 TensorRT 不支持改用支持的算子重写对应模块或用 OnnxGraphSurgeon 做子图替换处理这些问题的核心思路是不要硬改模型而是先复现最小问题再用等价算子组合替换或者调整导出参数。一个常见操作是用 onnx 库查看模型中的算子类型定位无法转换的节点。import onnx model onnx.load(bert_quantized.onnx) ops set() for node in model.graph.node: ops.add(node.op_type) print(ops)如果某个算子的输出与预期不符还可以用 onnxruntime 提供的调试工具检查图中间节点的输出逐一排查。4.3 校准数据对静态量化的影响一次真实的翻车经历想起一次典型的翻车经历给一个法律文本分类模型做静态量化用的是公开的通用分类数据做校准结果是 F1 直接暴跌 7 个点。一开始还怀疑是量化算法的问题后来换了 2000 条真实的法条文本做校准精度马上恢复到只降 0.2 个点。这件事让我养成了一个习惯校准数据集必须来自目标业务场景绝不能图省事用通用语料。这个教训写进了 Model-Optimizer 的流水线文档里当作强制要求。5. 模型优化的适用场景、边界与我的几点体会流程跑通之后我也更清楚这套流水线的适用边界。这里集中写几个容易被忽略但对实际决策很重要的点。5.1 什么时候适合用这套优化方案Model-Optimizer 这条流水线比较适合的典型场景包括模型已经训练完毕只想快速部署上线不希望重新训练。部署环境是 CPU 服务器或边缘设备资源受限。多个模型需要走同一套推理框架统一格式减少维护成本。团队里算法工程师和部署工程师之间需要协同统一的导出流程能减少沟通成本。如果你的场景是 GPU 部署TensorRT 会提供更多针对 GPU 的优化ONNX 导出通常还是作为中间步骤但优化方向会更多放在 FP16 精度和算子融合上量化到 INT8 的收益在 GPU 上没有 CPU 那么突出。如果你的模型结构有大量自定义层且每个层都没有 ONNX 对应算子那么走 ONNX 这条路就会很痛苦这时候可以先尝试把自定义层用标准算子重写再做下一步。5.2 关于量化是否导致精度严重损失的判断很多人问过一个问题量化之后精度一定会掉吗答案是不一定但大概率会有一点损失。影响损失大小的因素主要有三个模型本身的容量大模型冗余度高量化后精度损失通常较小小模型参数本身就少量化后更容易掉点。量化方式动态量化损失通常小于静态量化因为激活值没有被强制量化。校准数据的质量静态量化中校准数据分布和真实数据越接近损失越小。所以我会建议每个项目都给量化各步骤追加精度监控发现问题后优先调整校准数据集而不是急着换更复杂的量化算法。5.3 优化流程之外我还想分享的心得跑完这个项目后整理了几条经验先定指标再动手。上线前就把精度、时延、内存这三个数字写清楚优化过程中每一步都对着指标核算防止自己优化了个寂寞。保留每一步的中间产物。量化后的模型、导出前后的 ONNX 文件都保留一份出了问题可以快速对比定位。线程数、batch size 这些推理参数不要在项目中拍脑袋定。写一个简单的参数扫描脚本跑一个小网格通常几分钟就有结果但收益可能比换优化算法还大。遇到推理引擎不兼容问题优先考虑算子替换而不是放弃引擎。一个小算子替换往往就能让现有引擎继续正常工作。最后说一个我个人特别想强调的细节不要觉得某个优化手段效果好就一直用不同模型、不同部署环境适合的方案差别很大。Model-Optimizer 的价值不是某个单独技巧而是给了你一条可以在各个场景里快速调整复用的路径。量化、剪枝、蒸馏、导出、推理引擎调优每一步都是一个旋钮最佳配置永远是你在自己的数据、机器和业务指标下实测出来的而不是照抄任何人或任何文章的参数。我的建议是把自己项目的基线数据、每次调整后的对比结果好好记录下来这些数据积累到一定程度你就能形成一套属于自己团队的模型优化手感比任何通用的方法论都更可靠。
觉得有用,分享给同行:

为您的企业打造数字门面

稳重轻奢商务风格,端正雅致视觉,长效耐看不易过时。

立即咨询 →