PyTorch 2 Export 量化实战：借助 OpenVINOQuantizer 为 OpenVINO torch.compile 后端生成低精度模型

示例工程【免费下载链接】tutorialsPyTorch tutorials.项目地址https://gitcode.com/gh_mirrors/tuto/tutorials点击查看免费下载本篇技术指南完整讲解如何在 PyTorch 2 Export 量化流程中使用 Neural Network Compression FrameworkNNCF提供的OpenVINOQuantizer对 eager 模型依次执行 FX Graph 捕获、后训练量化PTQ、以及向 OpenVINO 表示的下沉lowering最终借助torch.compile(..., backendopenvino)在 Intel CPU 上获得可显著加速推理的量化模型。读完本文你将掌握torch.export捕获 ATen 图的完整链路、prepare_pt2e/convert_pt2e的量化三步走以及OpenVINOQuantizer的 preset、model_type、ignored_scope、target_device 等关键参数的调优方法。注意这是一项实验性功能量化 API 后续可能发生变化。本教程位于本仓库的unstable_source/openvino_quantizer.rst属于 PyTorch 的 prototype 教程目录相关特性作为技术预览提供PyTorch 官方不建议在生产流水线中使用详见 unstable_source/README.md。引言为什么需要 OpenVINOQuantizerPyTorch 2 Export 量化流程使用torch.export将模型捕获为计算图并在 ATen 图之上执行量化变换。相较于传统的量化方案这一思路预期能带来显著更高的模型覆盖率、更好的灵活性以及更简化的用户体验。而 OpenVINO 后端则负责把 TorchDynamo 生成的 FX Graph 编译为经过优化的 OpenVINO 模型。OpenVINOQuantizer来自 Neural Network Compression FrameworkNNCF其价值在于量化节点的放置位置是专门为 OpenVINO 设计的因此能够充分释放 OpenVINO 低精度low-precision算子的全部潜力。这一点正是本文所述流程的核心优势——量化方案与目标后端OpenVINO深度耦合而不是采用与后端无关的通用量化器。整体流程的四步骨架整个量化流程主要包含四个步骤Step 1基于 torch export 机制从 eager 模型捕获 FX Graph。Step 2基于捕获到的 FX Graph使用OpenVINOQuantizer执行 PyTorch 2 Export 量化流程。Step 3通过torch.compileAPI 将量化模型下沉lower为 OpenVINO 表示。可选 Step 4通过 NNCF 的quantize_pt2e方法进一步改善量化模型的精度指标。该流程的高层架构如下float_model(Python) Example Input \ / \ / —-------------------------------------------------------- | export | —-------------------------------------------------------- | FX Graph in ATen | | OpenVINOQuantizer | / —-------------------------------------------------------- | prepare_pt2e | | | | | Calibrate | | | | convert_pt2e | —-------------------------------------------------------- | Quantized Model | —-------------------------------------------------------- | Lower into Inductor | —-------------------------------------------------------- | OpenVINO model从图中可以清晰看到eager 模型与示例输入先经过export生成 ATen 层级的 FX Graph随后在OpenVINOQuantizer的指导下依次执行prepare_pt2e插入观察器、校准Calibrate、convert_pt2e转换为量化模型最后下沉到 Inductor得到最终的 OpenVINO 模型。环境准备安装 OpenVINO 与 NNCFOpenVINO 和 NNCF 可通过 pip 分发包轻松安装。在 Python 环境中依次执行pip install -U pip pip install openvino, nncf安装完成后即可在代码中导入import openvino.torch # 注册 openvino torch.compile 后端 import nncf.torch其中openvino.torch负责将 OpenVINO 后端注册到 PyTorch 的torch.compile后端列表中是后续backendopenvino能够生效的前提。第一步从 eager 模型捕获 FX Graph首先完成必要的导入并从 eager module 捕获 FX Graph。本教程以 torchvision 的 resnet18 模型为例进行后训练量化。import copy import openvino.torch import torch import torchvision.models as models from torch.ao.quantization.quantize_pt2e import convert_pt2e from torch.ao.quantization.quantize_pt2e import prepare_pt2e import nncf.torch # 创建 Eager 模型 model_name resnet18 model models.__dict__model_name # 将模型设置为 eval 模式 model model.eval() # 创建数据此处以 dummy 数据为例 traced_bs 50 x torch.randn(traced_bs, 3, 224, 224) example_inputs (x,) # 捕获待量化的 FX Graph with torch.no_grad(), nncf.torch.disable_patching(): exported_model torch.export.export(model, example_inputs).module()这里有几个值得注意的细节torch.export.export(model, example_inputs)是torch.export的主要入口。它以mod(*args, **kwargs)的调用方式追踪张量计算图并包装为一个ExportedProgram。该ExportedProgram可后续执行或序列化调用.module()可返回一个与原程序行为一致的、可调用的torch.nn.Module。ExportedProgram.graph属性即为我们所说的 FX Graph它包含 ATen IR即只包含 ATen 层级的运算graph_signature则详细描述了图中输入输出节点的类型参数、缓冲、用户输入、用户输出。这些细节可进一步参考本仓库的 intermediate_source/torch_export_tutorial.py其中注释详细说明了export()的签名export(mod, args, kwargsNone, *, dynamic_shapesNone) - ExportedProgram。nncf.torch.disable_patching()NNCF 通常会以 patch 的方式介入 PyTorch 内部实现例如为支持量化的算子做替换准备。在导出阶段将其禁用可以确保torch.export捕获到的是未被 NNCF 改动的、干净的原始计算图避免导出结果被 patch 干扰。torch.no_grad()导出和后续校准阶段都不需要梯度关闭自动求导可节省内存并避免不必要的图内容。为什么需要example_inputs导出过程依赖一组实际输入来追踪张量的形状与数据类型。这里的 batch size 为 50输入形状为(50, 3, 224, 224)即 50 张 224×224 的 RGB 图像与 ResNet18 的标准输入尺寸一致。如果后续要在真实推理中使用不同 batch size需要注意torch.export默认不支持下溢到不同形状需在导出时通过dynamic_shapes参数声明动态维度详见 intermediate_source/torch_export_tutorial.py 中对dynamic_shapes的说明。第二步使用 OpenVINOQuantizer 应用量化捕获到待量化的 FX Module 之后导入OpenVINOQuantizerfrom nncf.experimental.torch.fx import OpenVINOQuantizer quantizer OpenVINOQuantizer()OpenVINOQuantizer提供若干可选参数用于调质量化过程以获得精度更高的模型。下面列出关键参数及其含义preset定义模型的量化方案两种可用的 presetPERFORMANCE默认对权重和激活均使用对称量化。MIXED权重使用对称量化激活使用非对称量化。推荐用于包含 ELU、PReLU、GELU 等非 ReLU 或非对称激活函数的模型。OpenVINOQuantizer(presetnncf.QuantizationPreset.MIXED)model_type针对特定模型类型的专用量化方案用于指定特定类型模型所需的量化方案。目前唯一受支持的特殊量化方案是Transformer用于保持 BERT、Llama 等 Transformer 模型量化后的精度。默认值为None即不指定任何特殊方案。OpenVINOQuantizer(model_typenncf.ModelType.Transformer)ignored_scope将部分层排除在量化之外该参数可用于将某些层从量化过程中排除以保持模型精度。例如当希望排除模型的最后一层时可以按如下方式使用# 按层名排除 names [layer_1, layer_2, layer_3] OpenVINOQuantizer(ignored_scopenncf.IgnoredScope(namesnames)) # 按层类型排除 types [Conv2d, Linear] OpenVINOQuantizer(ignored_scopenncf.IgnoredScope(typestypes)) # 按正则表达式排除 regex .*layer_.* OpenVINOQuantizer(ignored_scopenncf.IgnoredScope(patternsregex)) # 按子图排除 # 此时图中从输入节点到输出节点的所有简单路径上的全部节点 # 都将被排除在量化过程之外。 subgraph nncf.Subgraph(inputs[layer_1, layer_2], outputs[layer_3]) OpenVINOQuantizer(ignored_scopenncf.IgnoredScope(subgraphs[subgraph]))这四种排除方式按名称、按类型、按正则、按子图分别适用于不同的精度修复场景名称排除最直观类型排除适合批量排除某类算子正则表达式适合有命名规律的层子图排除则适合排除一整条计算路径例如某个对精度敏感的分支。target_device指定优化所针对的目标设备定义目标设备优化时会考虑该设备的特性。支持以下取值ANY默认、CPU、CPU_SPR、GPU、NPU。OpenVINOQuantizer(target_devicenncf.TargetDevice.CPU)CPU_SPR指 Intel Sapphire Rapids第四代至强可扩展处理器其具备的高级矩阵扩展AMX能力可被 OpenVINO 编译器进一步利用NPU则面向 Intel 神经处理单元。选择与最终部署环境一致的设备可以让量化器生成的量化方案与该设备上可用的低精度算子集合对齐。执行 prepare → calibrate → convert 三步导入后端专用的 Quantizer 之后即可准备模型进行后训练量化。prepare_pt2e会执行两项关键工作将 BatchNorm 算子折叠fold到其前面的 Conv2d 算子中在模型中的合适位置插入观察器observers。prepared_model prepare_pt2e(exported_model, quantizer)插入观察器之后需要对prepared_model进行校准即用代表性数据跑一遍前向让观察器收集各张量实际的数据分布此处以 dummy 数据为例# 这里以 dummy 数据作为示例 prepared_model(*example_inputs)最后将校准后的模型转换为量化模型。convert_pt2e接收校准后的模型并产出量化模型quantized_model convert_pt2e(prepared_model, fold_quantizeFalse)fold_quantizeFalse的含义是在转换阶段不将量化/反量化算子折叠进权重中从而保留量化算子在图中的显式结构——这对于后续 OpenVINO 后端下沉阶段重新组织计算图、匹配低精度内核模式是有利的。完成以上步骤后量化流程即告结束得到量化模型quantized_model。第三步下沉Lower到 OpenVINO 表示得到量化模型后FX Graph 可以利用torch.compile(..., backendopenvino)功能获得 OpenVINO 的优化with torch.no_grad(), nncf.torch.disable_patching(): optimized_model torch.compile(quantized_model, backendopenvino) # 运行一些基准测试 optimized_model(*example_inputs)这里再次使用nncf.torch.disable_patching()包裹是为了确保下沉编译期间的计算图处理不受 NNCF patch 的干扰与第一步导出时的考虑一致。编译得到的optimized_model使用的是专为 Intel CPU 设计的低层级内核low-level kernels。与 eager 模型相比推理时间应当会显著缩短。需要说明的是文档作者在本教程中并未提供具体的基准数字实际加速幅度取决于模型结构、输入规模与硬件环境建议读者在自己的环境中以optimized_model(*example_inputs)这类方式自行测量对比。可选第四步通过 quantize_pt2e 改善量化模型指标NNCF 实现了多种高级量化算法例如SmoothQuant与BiasCorrection它们有助于在最小化原始模型与压缩模型输出差异的同时改善量化模型的精度指标。这些高级 NNCF 算法可以通过 NNCF 的quantize_pt2eAPI 访问from nncf.experimental.torch.fx import quantize_pt2e calibration_loader torch.utils.data.DataLoader(...) def transform_fn(data_item): images, _ data_item return images calibration_dataset nncf.Dataset(calibration_loader, transform_fn) quantized_model quantize_pt2e( exported_model, quantizer, calibration_dataset, smooth_quantTrue, fast_bias_correctionFalse )与手动的 prepare → calibrate → convert 流程相比quantize_pt2e是一个集成式的高阶入口calibration_dataset由nncf.Dataset包装的数据集transform_fn负责从 DataLoader 产出的数据项中提取模型实际需要的输入此处是从(images, labels)二元组中取出images。smooth_quantTrue启用 SmoothQuant 算法通过将激活的量化难度迁移到权重侧缓解激活值存在离群点时导致的量化精度损失对 Transformer 类模型尤其有效。fast_bias_correctionFalse关闭快速偏置校正模式使用更完整但更耗时的偏置校正流程通常在精度要求较高的场景下启用。该 API 内部会自动完成观察器的插入、基于数据集的校准以及最终的量化和高级算法的应用适合希望一键完成完整 PTQ 流水线的场景。结论与后续指引本教程介绍了如何将torch.compile与 OpenVINO 后端、以及 OpenVINO 专用量化器结合起来从 eager 模型出发经torch.export捕获 ATen FX Graph再由OpenVINOQuantizer配合prepare_pt2e/convert_pt2e完成后训练量化最终以torch.compile(..., backendopenvino)下沉为面向 Intel CPU 优化的 OpenVINO 模型。若需进一步改善精度可选用quantize_pt2e接入 SmoothQuant、BiasCorrection 等高级算法。相关仓库资源以下仓库内的文件可供进一步深入研读unstable_source/openvino_quantizer.rst本文所依据的原始教程文档包含完整的代码示例与参数说明。unstable_source/README.md说明本目录属于 prototype 教程与配方相关特性为技术预览不建议用于生产流水线。redirects.py其中记录了prototype/openvino_quantizer.html的重定向条目反映该教程曾位于 prototype 目录、现迁移至 unstable 目录的历史。intermediate_source/torch_export_tutorial.py深入讲解torch.export的入口签名、ExportedProgram的 graph / graph_signature 属性、ATen IR 语义以及 dynamic_shapes 动态形状用法是理解本文第一步捕获 FX Graph 原理的最佳配套材料。使用注意事项本流程依赖torch.export而torch.export不支持图断裂graph break若模型包含数据依赖的控制流等无法静态追踪的结构导出会失败需要先改写模型代码详见 intermediate_source/torch_export_tutorial.py 中关于 graph breaks 的讨论。本教程为 prototype 实验特性量化 API 可能在未来版本中变化升级 PyTorch、NNCF 或 OpenVINO 版本后需重新验证流程的兼容性。量化模型的实际推理加速与精度请以目标硬件CPU / CPU_SPR / GPU / NPU上的实测结果为准。赞分享示例工程【免费下载链接】tutorialsPyTorch tutorials.项目地址https://gitcode.com/gh_mirrors/tuto/tutorials点击查看免费下载相关推荐CANN graph-autofusion PyTorch 用例实战借助 torch.compile 与 ascendc 后端实现算子自动融合CANN graph autofusion PyTorch 用例实战借助 torch.compile 与 ascendc 后端实现算子自动融合 graph a人工智能模型编译AscendPyTorch Lightning Fabric 分布式低精度 Transformer 训练FP8 量化 FSDP2 张量并行 torch.compile 实战指南PyTorch Lightning Fabric 分布式低精度 Transformer 训练FP8 量化 FSDP2 张量并行 torch.compi人工智能深度学习机器学习预训练分布式训练微调PyTorch torch.compile 编程模型用 torch._dynamo.error_on_graph_break 精确控制图中断的行为PyTorch torch.compile 编程模型用 torch._dynamo.error_on_graph_break 精确控制图中断的行为 导读 t人工智能机器学习深度学习分布式训练模型编译上一篇Lucky Lillia Bot故障排查手册10个常见问题与终极解决方案下一篇EOS 钱包创建完全指南深入解析 cleos wallet create 命令的用法、原理与实战创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考