1. 模型优化器到底在解决什么问题第一次接触 Model-Optimizer 这个概念是在一个推荐系统的项目里。当时模型训练完离线指标 AUC 0.82看着挺漂亮一上线推理延迟直接飙到 800msQPS 连 50 都扛不住。老板问“能不能压到 100ms 以内”我盯着那坨 300MB 的 embedding 表心里只有一个念头这模型得“瘦身”。Model-Optimizer 说白了就是一套让模型跑得更快、占得更少、精度掉得可控的工具链和方法论。它不是某一个具体的库而是一个完整的优化体系涵盖量化、剪枝、蒸馏、算子融合、图优化、内存复用等多个维度。你可以在 PyTorch、TensorFlow、ONNX Runtime、TensorRT 这些框架里找到对应的实现但核心思想是相通的在不显著牺牲精度的前提下把计算量和内存占用压下来。什么人需要关注这个三类人最迫切。第一类是做端侧部署的手机、IoT 设备算力有限模型大一点就跑不动第二类是做在线服务的QPS 和延迟直接关系到服务器成本和用户体验第三类是做模型迭代的训练快一点、实验周期短一点论文产出就多一点。如果你只是跑跑 demo、做做学术实验那 Model-Optimizer 可能不是你的刚需但一旦涉及生产环境它就是绕不过去的坎。我见过太多团队在模型精度上死磕却忽略了推理效率这个“隐形成本”。一个 500MB 的模型和 50MB 的模型在云上的年费差距可能是六位数。所以这篇文章我想把 Model-Optimizer 的完整链路拆开从原理到实操从工具选型到踩坑记录尽量讲透。2. 优化方案的整体设计与选型逻辑2.1 先搞清楚瓶颈在哪计算密集还是内存密集优化最忌讳的就是“上来就量化”。我见过有人拿到模型直接上 INT8结果精度崩了 5 个点回头还得重新训。正确的做法是先做 profiling搞清楚瓶颈到底在哪。用 PyTorch Profiler 或者 Nsight Systems 跑一遍看时间花在哪些算子上。如果是 Conv、MatMul 这类计算密集型算子占大头那量化、TensorRT 加速收益会很明显如果是 Embedding lookup、LayerNorm 这类内存密集型算子占大头那量化收益有限反而要考虑内存布局优化、算子融合或者换更高效的实现。举个例子Transformer 类模型在推理时Attention 里的 QKV 计算是计算密集的但 KV Cache 的读写是内存密集的。长序列场景下KV Cache 可能占掉 70% 的显存。这时候你光量化权重没用得考虑 PagedAttention、KV Cache 量化或者 GQAGrouped Query Attention这类结构优化。实操心得profiling 至少跑三轮取稳定值。第一轮有冷启动开销第二轮有缓存预热第三轮才是真实表现。别偷懒只看一轮数据。2.2 量化、剪枝、蒸馏三条路怎么选量化是把 FP32 权重和激活值映射到 INT8 或 FP16直接减少内存带宽和计算量。剪枝是去掉不重要的权重或结构减少参数量。蒸馏是让小模型学大模型的行为换一个更小的架构。这三者不是互斥的实际项目里经常组合使用。我的经验是量化优先剪枝其次蒸馏最后。原因很简单量化对精度的影响通常最小尤其是训练后量化工程成本也最低剪枝需要精细调参搞不好精度掉得厉害蒸馏需要重新训练周期最长。但有个例外如果你的模型本身结构就冗余比如层数太多、通道数太大那剪枝的收益可能比量化更大。我做过一个实验一个 12 层的 BERT剪掉 4 层后精度只掉 0.3%但推理速度提升了 40%。这种情况下剪枝就是首选。2.3 训练后量化 vs 量化感知训练训练后量化PTQ是拿训练好的模型直接量化不需要重新训练工程成本低。量化感知训练QAT是在训练过程中模拟量化误差让模型适应低精度表示精度通常更好但需要重新训练。我的建议是先试 PTQ如果精度掉得在可接受范围内比如 1% 以内就直接用如果掉得太多再考虑 QAT。PTQ 里也有技巧比如用校准数据集做动态范围调整、逐通道量化、混合精度量化敏感层保持 FP16这些都能把精度拉回来不少。校准数据集的选择很关键。不要随便拿几百张图就完事要覆盖真实场景的分布。我做过一个图像分类的量化校准集用了 ImageNet 的 1000 张图结果某些类别的精度掉得特别厉害。后来发现是校准集里这些类别的样本太少动态范围估计不准。换成每类 50 张、总共 50000 张的校准集后精度就稳了。2.4 工具链选型PyTorch、ONNX、TensorRT 怎么配合PyTorch 生态里torch.quantization和torch.ao.quantization是官方量化工具支持 PTQ 和 QAT。ONNX Runtime 的量化工具更成熟支持多种量化模式而且跨框架。TensorRT 是 NVIDIA 的推理加速引擎量化算子融合内核自动调优性能最强但绑定 NVIDIA 硬件。我的典型工作流是PyTorch 训练 → 导出 ONNX → ONNX Runtime 量化 → TensorRT 部署。这样每一步都有回退空间ONNX 作为中间表示也方便调试。如果目标平台不是 NVIDIA那就用 ONNX Runtime 直接部署或者用 OpenVINOIntel、TFLite移动端。注意ONNX 导出时要注意算子版本和动态轴设置。我踩过坑动态轴设错了batch size 一变就报错。建议导出后用onnx.checker验证一遍再用onnxruntime跑个推理测试。3. 核心细节解析与实操要点3.1 量化参数怎么定从 min-max 到 KL 散度量化的核心是确定 scale 和 zero_point。最简单的是 min-max 校准直接取权重或激活的最小值和最大值作为范围。但这种方法对异常值敏感一个离群点就能把范围拉大导致大部分值量化后精度损失。更好的方法是 KL 散度校准通过最小化量化前后分布的 KL 散度来选范围。PyTorch 的torch.quantization.observer里就有HistogramObserver默认用 KL 散度。实测下来KL 散度比 min-max 在激活值量化上精度高 0.5-1 个点。逐通道量化per-channel比逐张量量化per-tensor精度更好尤其是卷积层。因为不同通道的权重分布差异很大逐通道能给每个通道单独算 scale。PyTorch 里设置torch.quantization.per_channel_weight_observer就行。代价是推理时多了一点计算开销但现代硬件上基本可以忽略。import torch.quantization as tq # 逐通道量化配置 qconfig tq.QConfig( activationtq.HistogramObserver.with_args( dtypetorch.quint8, qschemetorch.per_tensor_affine, reduce_rangeFalse ), weighttq.PerChannelMinMaxObserver.with_args( dtypetorch.qint8, qschemetorch.per_channel_symmetric ) )3.2 敏感层识别与混合精度策略不是所有层都适合量化。第一层和最后一层通常对精度影响最大因为第一层直接处理输入最后一层直接输出结果。Attention 里的 Softmax 层也敏感因为输出是概率分布量化误差会被放大。我的做法是先全量化跑一遍验证集看每层的量化误差。PyTorch 的torch.quantization提供了compare_weights和get_observer_dict可以逐层对比。误差大的层就保持 FP16 或 FP32其他层用 INT8。这就是混合精度量化。TensorRT 里可以用set_flag和set_precision逐层设置精度。ONNX Runtime 的QuantizationAwareTraining也支持op_types_to_quantize和nodes_to_exclude。实测下来混合精度量化通常能比全 INT8 提升 1-2 个点精度而速度只损失 5-10%。3.3 剪枝的粒度非结构化 vs 结构化非结构化剪枝是把单个权重置零稀疏度高但硬件加速难因为 GPU 对稀疏矩阵的支持有限。结构化剪枝是剪掉整个通道、整个头或者整个层硬件友好加速明显。我推荐结构化剪枝尤其是通道剪枝。用 L1 或 L2 范数衡量通道重要性剪掉范数最小的通道。但要注意剪枝后要 fine-tune 几轮让模型恢复精度。fine-tune 的学习率要小通常是原始学习率的 1/10 到 1/100。import torch.nn.utils.prune as prune # 对卷积层做 L1 结构化剪枝 module model.conv1 prune.ln_structured(module, nameweight, amount0.3, n1, dim0) # 剪枝后要 fine-tune实操心得剪枝比例不要一次剪太多建议迭代式剪枝。先剪 10%fine-tune 恢复再剪 10%再 fine-tune。这样最终能剪掉 50% 甚至更多而精度只掉 1-2 个点。一次性剪 50%精度基本就崩了。3.4 算子融合与图优化算子融合是把多个小算子合并成一个大算子减少内核启动开销和内存读写。比如 Conv BN ReLU 融合成一个算子这是最常见的。TensorRT 和 ONNX Runtime 都会自动做这种融合但前提是图结构要清晰。手动融合的话PyTorch 里可以用torch.quantization.fuse_modules。但要注意融合后的算子量化行为可能和单独量化不同需要重新校准。# 融合 Conv BN ReLU fused_model torch.quantization.fuse_modules( model, [[conv1, bn1, relu1]], inplaceFalse )图优化还包括常量折叠、死代码消除、布局转换等。ONNX Runtime 的graph_optimization_level可以设置优化级别ORT_ENABLE_ALL会开启所有优化。但有时候过度优化会导致精度问题建议先开ORT_ENABLE_BASIC验证没问题再逐步提高。4. 完整实操流程与关键环节实现4.1 环境准备与依赖安装先建一个干净的虚拟环境避免依赖冲突。Python 3.8 以上PyTorch 1.12 以上ONNX Runtime 1.14 以上TensorRT 8.5 以上。如果要用 GPU 加速CUDA 版本要和 TensorRT 匹配。conda create -n model-opt python3.9 conda activate model-opt pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install onnx onnxruntime-gpu pip install tensorrt pip install polygraphy # TensorRT 调试工具注意TensorRT 的安装比较麻烦建议用 NVIDIA 官方提供的 tar 包或者 Docker 镜像。pip 安装的版本可能和 CUDA 不匹配跑起来会报libcudnn.so找不到。4.2 基线模型导出与验证拿一个训练好的 ResNet50 做例子。先导出 ONNX设置动态 batch size。import torch import torchvision.models as models model models.resnet50(pretrainedTrue) model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, resnet50.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}}, opset_version13 )导出后用 ONNX Runtime 跑一遍对比 PyTorch 的输出确保误差在 1e-5 以内。import onnxruntime as ort import numpy as np sess ort.InferenceSession(resnet50.onnx) input_data np.random.randn(1, 3, 224, 224).astype(np.float32) onnx_output sess.run(None, {input: input_data})[0] torch_output model(torch.from_numpy(input_data)).detach().numpy() print(Max diff:, np.max(np.abs(onnx_output - torch_output)))4.3 训练后量化实操用 ONNX Runtime 的量化工具做 PTQ。先准备校准数据集这里用 ImageNet 的验证集取 1000 张。from onnxruntime.quantization import quantize_static, CalibrationDataReader, QuantType class ImageNetCalibrationReader(CalibrationDataReader): def __init__(self, calibration_data): self.data calibration_data self.index 0 def get_next(self): if self.index len(self.data): return None batch self.data[self.index] self.index 1 return {input: batch} calibration_data [...] # 预处理好的 numpy 数组列表 reader ImageNetCalibrationReader(calibration_data) quantize_static( model_inputresnet50.onnx, model_outputresnet50_quant.onnx, calibration_data_readerreader, quant_formatQuantType.QInt8, per_channelTrue, reduce_rangeFalse )量化完再跑一遍精度对比。如果掉点超过 1%就调整校准集或者改用混合精度。4.4 TensorRT 引擎构建与推理TensorRT 的流程是ONNX → 解析 → 优化 → 序列化引擎 → 推理。用trtexec命令行工具最方便。trtexec --onnxresnet50_quant.onnx \ --saveEngineresnet50_quant.engine \ --fp16 \ --int8 \ --workspace4096 \ --verbose--fp16和--int8同时开TensorRT 会自动选择每层的最优精度。--workspace是显存上限单位 MB。--verbose打印详细日志方便排查哪层用了什么精度。推理时用 Python APIimport tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit logger trt.Logger(trt.Logger.WARNING) with open(resnet50_quant.engine, rb) as f: engine trt.Runtime(logger).deserialize_cuda_engine(f.read()) context engine.create_execution_context() # 分配输入输出显存执行推理实操心得TensorRT 引擎和硬件绑定换 GPU 型号要重新构建。构建过程可能很慢ResNet50 大概 2-3 分钟BERT 可能 10 分钟以上。建议把构建好的引擎存下来部署时直接加载。4.5 性能对比与精度验证跑一个完整的对比表格包括延迟、吞吐、显存占用、精度。模型版本延迟 (ms)吞吐 (QPS)显存 (MB)Top-1 精度PyTorch FP324522120076.1%ONNX FP323826110076.1%ONNX INT8185560075.3%TensorRT FP16128350076.0%TensorRT INT8812540075.1%从数据看TensorRT INT8 比 PyTorch FP32 快了 5 倍多精度只掉 1 个点。这个收益在生产环境里非常可观。5. 常见问题与排查技巧实录5.1 量化后精度暴跌怎么办先定位是哪一层的问题。用onnxruntime.quantization的quantize_static时可以设置nodes_to_exclude排除敏感层。或者用QuantFormat.QDQ模式插入 QuantizeLinear/DequantizeLinear 节点方便逐层对比。如果精度还是不行就上 QAT。PyTorch 的 QAT 流程是插入伪量化节点 → 微调 → 转换为量化模型。微调通常 5-10 个 epoch 就够了学习率用 1e-5 到 1e-4。model.qconfig torch.quantization.get_default_qat_qconfig(fbgemm) model_prepared torch.quantization.prepare_qat(model, inplaceFalse) # 微调 model_prepared.train() for epoch in range(10): for data, target in train_loader: optimizer.zero_grad() output model_prepared(data) loss criterion(output, target) loss.backward() optimizer.step() model_quantized torch.quantization.convert(model_prepared.eval(), inplaceFalse)5.2 TensorRT 构建失败常见原因最常见的是算子不支持。TensorRT 不是所有 ONNX 算子都支持尤其是一些自定义算子或者新版本的算子。用polygraphy可以查看哪些算子不支持。polygraphy run resnet50.onnx --trt --verbose如果是不支持的算子要么用 TensorRT 的 Plugin 自己实现要么在 ONNX 里替换成支持的算子组合。另一个常见原因是动态 shape 设置不对TensorRT 需要显式指定 optimization profile。profile builder.create_optimization_profile() profile.set_shape(input, min(1, 3, 224, 224), opt(8, 3, 224, 224), max(32, 3, 224, 224)) config.add_optimization_profile(profile)5.3 显存不够用怎么优化除了量化还有几个技巧。第一用torch.cuda.empty_cache()及时释放缓存。第二用梯度检查点gradient checkpointing换显存训练时用推理时不用。第三用pin_memory和num_workers优化数据加载避免数据加载成为瓶颈。推理时如果 batch size 太大导致显存不够可以用动态 batch 或者流式推理。TensorRT 的 optimization profile 里设置 max batch运行时根据实际显存调整。5.4 常见问题速查表问题现象可能原因排查方法解决方案量化后精度掉 3%校准集分布不对对比校准集和验证集分布换校准集增加样本量TensorRT 构建报错算子不支持polygraphy run --trt替换算子或写 Plugin推理延迟波动大动态 shape 未优化查看 profile 设置设置合理的 opt shape显存 OOMbatch size 太大nvidia-smi监控减小 batch 或量化ONNX 导出失败算子版本不匹配onnx.checker验证调整 opset_version量化模型推理慢未融合算子查看图优化日志开启ORT_ENABLE_ALL最后再分享一个小技巧量化校准的时候可以用reduce_rangeTrue来避免 INT8 溢出。这个参数会把量化范围从 [-128, 127] 缩到 [-64, 63]牺牲一点动态范围换稳定性。在激活值分布比较集中的模型上效果很好。这个内容后续还可以这样扩展针对特定硬件平台如高通 Hexagon、苹果 Neural Engine做定制优化或者结合 Neural Architecture Search 做自动化模型压缩。但那是另一个话题了有机会再聊。 Model-Optimizer实战：量化剪枝与TensorRT推理加速全链路
1. 模型优化器到底在解决什么问题第一次接触 Model-Optimizer 这个概念是在一个推荐系统的项目里。当时模型训练完离线指标 AUC 0.82看着挺漂亮一上线推理延迟直接飙到 800msQPS 连 50 都扛不住。老板问“能不能压到 100ms 以内”我盯着那坨 300MB 的 embedding 表心里只有一个念头这模型得“瘦身”。Model-Optimizer 说白了就是一套让模型跑得更快、占得更少、精度掉得可控的工具链和方法论。它不是某一个具体的库而是一个完整的优化体系涵盖量化、剪枝、蒸馏、算子融合、图优化、内存复用等多个维度。你可以在 PyTorch、TensorFlow、ONNX Runtime、TensorRT 这些框架里找到对应的实现但核心思想是相通的在不显著牺牲精度的前提下把计算量和内存占用压下来。什么人需要关注这个三类人最迫切。第一类是做端侧部署的手机、IoT 设备算力有限模型大一点就跑不动第二类是做在线服务的QPS 和延迟直接关系到服务器成本和用户体验第三类是做模型迭代的训练快一点、实验周期短一点论文产出就多一点。如果你只是跑跑 demo、做做学术实验那 Model-Optimizer 可能不是你的刚需但一旦涉及生产环境它就是绕不过去的坎。我见过太多团队在模型精度上死磕却忽略了推理效率这个“隐形成本”。一个 500MB 的模型和 50MB 的模型在云上的年费差距可能是六位数。所以这篇文章我想把 Model-Optimizer 的完整链路拆开从原理到实操从工具选型到踩坑记录尽量讲透。2. 优化方案的整体设计与选型逻辑2.1 先搞清楚瓶颈在哪计算密集还是内存密集优化最忌讳的就是“上来就量化”。我见过有人拿到模型直接上 INT8结果精度崩了 5 个点回头还得重新训。正确的做法是先做 profiling搞清楚瓶颈到底在哪。用 PyTorch Profiler 或者 Nsight Systems 跑一遍看时间花在哪些算子上。如果是 Conv、MatMul 这类计算密集型算子占大头那量化、TensorRT 加速收益会很明显如果是 Embedding lookup、LayerNorm 这类内存密集型算子占大头那量化收益有限反而要考虑内存布局优化、算子融合或者换更高效的实现。举个例子Transformer 类模型在推理时Attention 里的 QKV 计算是计算密集的但 KV Cache 的读写是内存密集的。长序列场景下KV Cache 可能占掉 70% 的显存。这时候你光量化权重没用得考虑 PagedAttention、KV Cache 量化或者 GQAGrouped Query Attention这类结构优化。实操心得profiling 至少跑三轮取稳定值。第一轮有冷启动开销第二轮有缓存预热第三轮才是真实表现。别偷懒只看一轮数据。2.2 量化、剪枝、蒸馏三条路怎么选量化是把 FP32 权重和激活值映射到 INT8 或 FP16直接减少内存带宽和计算量。剪枝是去掉不重要的权重或结构减少参数量。蒸馏是让小模型学大模型的行为换一个更小的架构。这三者不是互斥的实际项目里经常组合使用。我的经验是量化优先剪枝其次蒸馏最后。原因很简单量化对精度的影响通常最小尤其是训练后量化工程成本也最低剪枝需要精细调参搞不好精度掉得厉害蒸馏需要重新训练周期最长。但有个例外如果你的模型本身结构就冗余比如层数太多、通道数太大那剪枝的收益可能比量化更大。我做过一个实验一个 12 层的 BERT剪掉 4 层后精度只掉 0.3%但推理速度提升了 40%。这种情况下剪枝就是首选。2.3 训练后量化 vs 量化感知训练训练后量化PTQ是拿训练好的模型直接量化不需要重新训练工程成本低。量化感知训练QAT是在训练过程中模拟量化误差让模型适应低精度表示精度通常更好但需要重新训练。我的建议是先试 PTQ如果精度掉得在可接受范围内比如 1% 以内就直接用如果掉得太多再考虑 QAT。PTQ 里也有技巧比如用校准数据集做动态范围调整、逐通道量化、混合精度量化敏感层保持 FP16这些都能把精度拉回来不少。校准数据集的选择很关键。不要随便拿几百张图就完事要覆盖真实场景的分布。我做过一个图像分类的量化校准集用了 ImageNet 的 1000 张图结果某些类别的精度掉得特别厉害。后来发现是校准集里这些类别的样本太少动态范围估计不准。换成每类 50 张、总共 50000 张的校准集后精度就稳了。2.4 工具链选型PyTorch、ONNX、TensorRT 怎么配合PyTorch 生态里torch.quantization和torch.ao.quantization是官方量化工具支持 PTQ 和 QAT。ONNX Runtime 的量化工具更成熟支持多种量化模式而且跨框架。TensorRT 是 NVIDIA 的推理加速引擎量化算子融合内核自动调优性能最强但绑定 NVIDIA 硬件。我的典型工作流是PyTorch 训练 → 导出 ONNX → ONNX Runtime 量化 → TensorRT 部署。这样每一步都有回退空间ONNX 作为中间表示也方便调试。如果目标平台不是 NVIDIA那就用 ONNX Runtime 直接部署或者用 OpenVINOIntel、TFLite移动端。注意ONNX 导出时要注意算子版本和动态轴设置。我踩过坑动态轴设错了batch size 一变就报错。建议导出后用onnx.checker验证一遍再用onnxruntime跑个推理测试。3. 核心细节解析与实操要点3.1 量化参数怎么定从 min-max 到 KL 散度量化的核心是确定 scale 和 zero_point。最简单的是 min-max 校准直接取权重或激活的最小值和最大值作为范围。但这种方法对异常值敏感一个离群点就能把范围拉大导致大部分值量化后精度损失。更好的方法是 KL 散度校准通过最小化量化前后分布的 KL 散度来选范围。PyTorch 的torch.quantization.observer里就有HistogramObserver默认用 KL 散度。实测下来KL 散度比 min-max 在激活值量化上精度高 0.5-1 个点。逐通道量化per-channel比逐张量量化per-tensor精度更好尤其是卷积层。因为不同通道的权重分布差异很大逐通道能给每个通道单独算 scale。PyTorch 里设置torch.quantization.per_channel_weight_observer就行。代价是推理时多了一点计算开销但现代硬件上基本可以忽略。import torch.quantization as tq # 逐通道量化配置 qconfig tq.QConfig( activationtq.HistogramObserver.with_args( dtypetorch.quint8, qschemetorch.per_tensor_affine, reduce_rangeFalse ), weighttq.PerChannelMinMaxObserver.with_args( dtypetorch.qint8, qschemetorch.per_channel_symmetric ) )3.2 敏感层识别与混合精度策略不是所有层都适合量化。第一层和最后一层通常对精度影响最大因为第一层直接处理输入最后一层直接输出结果。Attention 里的 Softmax 层也敏感因为输出是概率分布量化误差会被放大。我的做法是先全量化跑一遍验证集看每层的量化误差。PyTorch 的torch.quantization提供了compare_weights和get_observer_dict可以逐层对比。误差大的层就保持 FP16 或 FP32其他层用 INT8。这就是混合精度量化。TensorRT 里可以用set_flag和set_precision逐层设置精度。ONNX Runtime 的QuantizationAwareTraining也支持op_types_to_quantize和nodes_to_exclude。实测下来混合精度量化通常能比全 INT8 提升 1-2 个点精度而速度只损失 5-10%。3.3 剪枝的粒度非结构化 vs 结构化非结构化剪枝是把单个权重置零稀疏度高但硬件加速难因为 GPU 对稀疏矩阵的支持有限。结构化剪枝是剪掉整个通道、整个头或者整个层硬件友好加速明显。我推荐结构化剪枝尤其是通道剪枝。用 L1 或 L2 范数衡量通道重要性剪掉范数最小的通道。但要注意剪枝后要 fine-tune 几轮让模型恢复精度。fine-tune 的学习率要小通常是原始学习率的 1/10 到 1/100。import torch.nn.utils.prune as prune # 对卷积层做 L1 结构化剪枝 module model.conv1 prune.ln_structured(module, nameweight, amount0.3, n1, dim0) # 剪枝后要 fine-tune实操心得剪枝比例不要一次剪太多建议迭代式剪枝。先剪 10%fine-tune 恢复再剪 10%再 fine-tune。这样最终能剪掉 50% 甚至更多而精度只掉 1-2 个点。一次性剪 50%精度基本就崩了。3.4 算子融合与图优化算子融合是把多个小算子合并成一个大算子减少内核启动开销和内存读写。比如 Conv BN ReLU 融合成一个算子这是最常见的。TensorRT 和 ONNX Runtime 都会自动做这种融合但前提是图结构要清晰。手动融合的话PyTorch 里可以用torch.quantization.fuse_modules。但要注意融合后的算子量化行为可能和单独量化不同需要重新校准。# 融合 Conv BN ReLU fused_model torch.quantization.fuse_modules( model, [[conv1, bn1, relu1]], inplaceFalse )图优化还包括常量折叠、死代码消除、布局转换等。ONNX Runtime 的graph_optimization_level可以设置优化级别ORT_ENABLE_ALL会开启所有优化。但有时候过度优化会导致精度问题建议先开ORT_ENABLE_BASIC验证没问题再逐步提高。4. 完整实操流程与关键环节实现4.1 环境准备与依赖安装先建一个干净的虚拟环境避免依赖冲突。Python 3.8 以上PyTorch 1.12 以上ONNX Runtime 1.14 以上TensorRT 8.5 以上。如果要用 GPU 加速CUDA 版本要和 TensorRT 匹配。conda create -n model-opt python3.9 conda activate model-opt pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install onnx onnxruntime-gpu pip install tensorrt pip install polygraphy # TensorRT 调试工具注意TensorRT 的安装比较麻烦建议用 NVIDIA 官方提供的 tar 包或者 Docker 镜像。pip 安装的版本可能和 CUDA 不匹配跑起来会报libcudnn.so找不到。4.2 基线模型导出与验证拿一个训练好的 ResNet50 做例子。先导出 ONNX设置动态 batch size。import torch import torchvision.models as models model models.resnet50(pretrainedTrue) model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, resnet50.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}}, opset_version13 )导出后用 ONNX Runtime 跑一遍对比 PyTorch 的输出确保误差在 1e-5 以内。import onnxruntime as ort import numpy as np sess ort.InferenceSession(resnet50.onnx) input_data np.random.randn(1, 3, 224, 224).astype(np.float32) onnx_output sess.run(None, {input: input_data})[0] torch_output model(torch.from_numpy(input_data)).detach().numpy() print(Max diff:, np.max(np.abs(onnx_output - torch_output)))4.3 训练后量化实操用 ONNX Runtime 的量化工具做 PTQ。先准备校准数据集这里用 ImageNet 的验证集取 1000 张。from onnxruntime.quantization import quantize_static, CalibrationDataReader, QuantType class ImageNetCalibrationReader(CalibrationDataReader): def __init__(self, calibration_data): self.data calibration_data self.index 0 def get_next(self): if self.index len(self.data): return None batch self.data[self.index] self.index 1 return {input: batch} calibration_data [...] # 预处理好的 numpy 数组列表 reader ImageNetCalibrationReader(calibration_data) quantize_static( model_inputresnet50.onnx, model_outputresnet50_quant.onnx, calibration_data_readerreader, quant_formatQuantType.QInt8, per_channelTrue, reduce_rangeFalse )量化完再跑一遍精度对比。如果掉点超过 1%就调整校准集或者改用混合精度。4.4 TensorRT 引擎构建与推理TensorRT 的流程是ONNX → 解析 → 优化 → 序列化引擎 → 推理。用trtexec命令行工具最方便。trtexec --onnxresnet50_quant.onnx \ --saveEngineresnet50_quant.engine \ --fp16 \ --int8 \ --workspace4096 \ --verbose--fp16和--int8同时开TensorRT 会自动选择每层的最优精度。--workspace是显存上限单位 MB。--verbose打印详细日志方便排查哪层用了什么精度。推理时用 Python APIimport tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit logger trt.Logger(trt.Logger.WARNING) with open(resnet50_quant.engine, rb) as f: engine trt.Runtime(logger).deserialize_cuda_engine(f.read()) context engine.create_execution_context() # 分配输入输出显存执行推理实操心得TensorRT 引擎和硬件绑定换 GPU 型号要重新构建。构建过程可能很慢ResNet50 大概 2-3 分钟BERT 可能 10 分钟以上。建议把构建好的引擎存下来部署时直接加载。4.5 性能对比与精度验证跑一个完整的对比表格包括延迟、吞吐、显存占用、精度。模型版本延迟 (ms)吞吐 (QPS)显存 (MB)Top-1 精度PyTorch FP324522120076.1%ONNX FP323826110076.1%ONNX INT8185560075.3%TensorRT FP16128350076.0%TensorRT INT8812540075.1%从数据看TensorRT INT8 比 PyTorch FP32 快了 5 倍多精度只掉 1 个点。这个收益在生产环境里非常可观。5. 常见问题与排查技巧实录5.1 量化后精度暴跌怎么办先定位是哪一层的问题。用onnxruntime.quantization的quantize_static时可以设置nodes_to_exclude排除敏感层。或者用QuantFormat.QDQ模式插入 QuantizeLinear/DequantizeLinear 节点方便逐层对比。如果精度还是不行就上 QAT。PyTorch 的 QAT 流程是插入伪量化节点 → 微调 → 转换为量化模型。微调通常 5-10 个 epoch 就够了学习率用 1e-5 到 1e-4。model.qconfig torch.quantization.get_default_qat_qconfig(fbgemm) model_prepared torch.quantization.prepare_qat(model, inplaceFalse) # 微调 model_prepared.train() for epoch in range(10): for data, target in train_loader: optimizer.zero_grad() output model_prepared(data) loss criterion(output, target) loss.backward() optimizer.step() model_quantized torch.quantization.convert(model_prepared.eval(), inplaceFalse)5.2 TensorRT 构建失败常见原因最常见的是算子不支持。TensorRT 不是所有 ONNX 算子都支持尤其是一些自定义算子或者新版本的算子。用polygraphy可以查看哪些算子不支持。polygraphy run resnet50.onnx --trt --verbose如果是不支持的算子要么用 TensorRT 的 Plugin 自己实现要么在 ONNX 里替换成支持的算子组合。另一个常见原因是动态 shape 设置不对TensorRT 需要显式指定 optimization profile。profile builder.create_optimization_profile() profile.set_shape(input, min(1, 3, 224, 224), opt(8, 3, 224, 224), max(32, 3, 224, 224)) config.add_optimization_profile(profile)5.3 显存不够用怎么优化除了量化还有几个技巧。第一用torch.cuda.empty_cache()及时释放缓存。第二用梯度检查点gradient checkpointing换显存训练时用推理时不用。第三用pin_memory和num_workers优化数据加载避免数据加载成为瓶颈。推理时如果 batch size 太大导致显存不够可以用动态 batch 或者流式推理。TensorRT 的 optimization profile 里设置 max batch运行时根据实际显存调整。5.4 常见问题速查表问题现象可能原因排查方法解决方案量化后精度掉 3%校准集分布不对对比校准集和验证集分布换校准集增加样本量TensorRT 构建报错算子不支持polygraphy run --trt替换算子或写 Plugin推理延迟波动大动态 shape 未优化查看 profile 设置设置合理的 opt shape显存 OOMbatch size 太大nvidia-smi监控减小 batch 或量化ONNX 导出失败算子版本不匹配onnx.checker验证调整 opset_version量化模型推理慢未融合算子查看图优化日志开启ORT_ENABLE_ALL最后再分享一个小技巧量化校准的时候可以用reduce_rangeTrue来避免 INT8 溢出。这个参数会把量化范围从 [-128, 127] 缩到 [-64, 63]牺牲一点动态范围换稳定性。在激活值分布比较集中的模型上效果很好。这个内容后续还可以这样扩展针对特定硬件平台如高通 Hexagon、苹果 Neural Engine做定制优化或者结合 Neural Architecture Search 做自动化模型压缩。但那是另一个话题了有机会再聊。