TensorRT 量化工具包自定义 Q/DQ 插入详解：用 CustomQDQInsertionCase 实现 ResNet 残差连接的 INT8 最优融合

TensorRT 量化工具包自定义 Q/DQ 插入详解用 CustomQDQInsertionCase 实现 ResNet 残差连接的 INT8 最优融合【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRTTensorFlow 量化工具包tensorflow-quantization本仓库位于 tools/tensorflow-quantization默认按 NVIDIA 量化配方为受支持层插入 Q/DQ 节点但默认行为并不总能产生对 TensorRT INT8 内核融合最有利的图结构。本文以 ResNet 残差连接为实战场景系统讲解自定义 Q/DQ 插入Custom Q/DQ Insertion Case特性的设计动机、库内置 Cases 清单以及如何通过继承CustomQDQInsertionCase类编写自己的 Q/DQ 插入策略最终让量化后的 ONNX 模型在 TensorRT 中触发最优的 ConvAdd 融合。读完本文你将掌握quantize_model的custom_qdq_cases参数、QuantizationSpec与CustomQDQInsertionCase的分工以及一套可复用的自定义 Case 开发范式。为什么需要自定义 Q/DQ 插入工具包对每个受支持层的默认量化行为展示在 Add New Layer Support 一节的支持层矩阵中。总体而言它会对受支持层的所有输入以及所有权重若该层有权重插入 Q/DQ 节点。但默认行为未必总能带来最优的 TensorRT INT8 融合例如在 ResNet 这类带残差连接的模型中需要在残差连接上添加 Q/DQ 节点而默认行为恰恰不会在Add层前插入 Q/DQ若不加区分地对Add层所有输入都量化又会在Conv → Add之间形成多余的 Q/DQ 对阻碍Conv ElementWise(Add)融合成单个卷积内核。为了应对这类网络特定需求工具包提供了Custom Q/DQ Insertion Case库特性允许用户以编程方式决定某个特定层在特定场景下应该如何被差异化量化。需要指出的是直接提供一个QuantizationSpec类对象是实现同一目标的硬编码方式能改但不够智能而自定义 Q/DQ Case 则是更灵活、可复用的方式。从源码看该特性的入口是 quantize.py 中的quantize_model函数其完整签名为def quantize_model( model, quantization_mode: str full, quantization_spec: QuantizationSpec None, custom_qdq_cases: List[CustomQDQInsertionCase] None, ) - tf.keras.Model在函数内部见 quantize.py工具包会依次调用custom_qdq_cases列表中每个对象的case(model, quantization_spec)方法收集返回的QuantizationSpec再与用户传入的quantization_spec合并后进入实际量化流程。这正是自定义 Case 影响 Q/DQ 插入的底层机制。实战背景以 ResNet 残差连接为例的三种量化方案文档用一个 ResNet 风格模型作为贯穿案例目标是对模型的残差连接执行自定义量化。下面依次展示三种量化场景默认量化、基于QuantizationSpec的自定义量化次优、基于Custom Q/DQ Insertion Case的自定义量化最优。场景一默认量化Default Quantization默认量化对应基础文档中的Full Default Quantization模式见 basics.md即所有受支持层按默认行为量化代码只需一行# Quantize model q_nn_model quantize_model(modelnn_model_original)下图展示了基线 ResNet 残差块左与默认量化后的残差块右。可以看到默认行为不会在Add层前插入 Q/DQ 节点。图 1. ResNet 残差块左与默认量化后的残差块右Add 层未被量化需要说明的是AddQuantizeWrapper在工具包中已经实现只是默认被禁用对应 add_new_layer_support.md 支持矩阵中tf.keras.layers.Add的Quantize Input False*。因此量化Add层最简单的方式就是启用对类类型为Add的层的量化。场景二基于 QuantizationSpec 的自定义量化次优该方案对应Full Custom Quantization模式通过QuantizationSpec启用所有Add类的量化# 1. Enable Add layer quantization qspec QuantizationSpec() qspec.add(nameAdd, is_keras_classTrue) # 2. Quantize model q_nn_model quantize_model( modelnn_model_original, quantization_specqspec )QuantizationSpec的核心是add()方法其参数定义于 quantize.py包括参数类型默认值含义namestr/List必填层名model.summary()中可见或 Keras 类名配合is_keras_classTrueis_keras_classbool/ListFalse为True时name代表 Keras 层类quantize_inputbool/ListTrue是否量化该层的输入默认按 NVIDIA 配方quantize_weightbool/ListTrue是否量化该层的权重无权重层忽略此值quantization_indexList/List[List]None对多输入层如Add、Concatenate指定量化哪些输入索引下图展示了该方案下的量化结果Add层的所有输入都被量化了。图 2. 通过 QuantizationSpec 对 Add 层量化右其所有输入都被插入 Q/DQ为什么是次优虽然Add被量化了但在Conv → Add这条路径上多出的 Q/DQ 对依然无法触发 TensorRT 中卷积层后接 ElementWise如Add层融合为单个卷积内核的最优融合。TensorRT 的 Q/DQ 放置建议是只在残差连接分支上添加 Q/DQ 节点而不是在Add与Conv之间。场景三基于 Custom Q/DQ Insertion Case 的自定义量化最优库提供的ResNetV1QDQCase类解决了上述问题它通过编程让Add层在某一路径连接到Conv时跳过该路径的 Q/DQ 插入。使用方法是在quantize_model中传入 Case 对象列表# 1. Indicate one or more custom QDQ cases custom_qdq_case ResNetV1QDQCase() # 3. Quantize model q_nn_model quantize_model( modelnn_model_original, custom_qdq_cases[custom_qdq_case] )图 3. 通过 ResNetV1QDQCase 量化右只有残差分支被插入 Q/DQ对比图 2 与图 3 可以观察到两个关键差异Conv路径不再被插入 Q/DQ来自Conv层的输入路径保持浮点直通仅残差分支被量化从而为 TensorRT 的ConvAdd融合保留了干净的图结构水平融合Horizontal Fusion由于第一个Relu层的两个输出都被量化两处 Q/DQ 对被合并为一对进一步精简了图的节点数。这种量化方式得到的图对 TensorRT INT8 融合而言是最优的。库内置的 Custom Q/DQ Insertion Cases工具包为模型库model zoo中已有的模型提供了现成的自定义 Q/DQ 插入 Case。所有类都可以从tensorflow_quantization.custom_qdq_cases模块导入并作为对象传给quantize_model函数。完整的模块实现在 custom_qdq_cases.py。Case 类目标模型核心作用ResidualConnectionQDQCase通用残差网络对Add层的残差输入插入 Q/DQ当输入为 Conv / Conv-BN / Conv-BN-Activation / Conv-BN-Activation-Dropout 模式且只有一个输出分支时跳过该分支只量化其余输入以触发 Add 的 INT8 融合BNQDQCaseResNet-v2 等当 BatchNorm 连接到 Conv 层时跳过其 Q/DQ该 BN 会被前序 Conv 融合仅量化不与 Conv 相连的 BN如 pre-activation 结构中的 BNMaxPoolQDQCaseResNet-v2、Inception 等启用 MaxPooling2D 层的输入量化用于残差连接/Concat 输出处的水平融合ResNetV1QDQCaseResNet-v1组合ResidualConnectionQDQCase量化全部 ResNet-v1 残差连接ResNetV2QDQCaseResNet-v2组合ResidualConnectionQDQCaseBNQDQCaseMaxPoolQDQCaseEfficientNetQDQCaseEfficientNet-B0 等残差连接 SESqueeze-Excite块中Multiply层输入索引 0、1 的量化MobileNetQDQCaseMobileNet-v1/v2残差连接量化InceptionQDQCaseInception-v3MaxPoolQDQCase在 Concat 输出处触发水平融合从源码结构看这些 Case 类体现了可组合设计例如ResNetV2QDQCase.case内部依次调用ResidualConnectionQDQCase、BNQDQCase、MaxPoolQDQCase的case方法并将各自返回的QuantizationSpec.layers拼接为一个新的QuantizationSpec见 custom_qdq_cases.py。这意味着你在为自己的网络编写 Case 时也可以直接复用这些通用 Case 作为积木。from tensorflow_quantization.custom_qdq_cases import ResNetV1QDQCase from tensorflow_quantization import quantize_model q_model quantize_model(modelresnet50, custom_qdq_cases[ResNetV1QDQCase()])这些内置 Case 的行为有测试用例背书见 tests/custom_qdq_cases_test.py例如test_resnet50_residual_qdq_case断言 ResNet-50 各Add层的quantization_index满足预期残差块首层Add与带投影的Add按规则处理test_resnet50v2_bn_qdq_case则验证了 ResNet-50V2 中preact_bn、post_bn等 BatchNorm 层与 MaxPool 层被正确捕获。如何添加一个新的 Custom Q/DQ Insertion Case第一步继承基类新 Case 类必须继承tensorflow_quantization.CustomQDQInsertionCase基类。该基类定义于 custom_qdq_case_base.py是一个ABC抽象类包含两个方法class CustomQDQInsertionCase(ABC): def info(self) - str: return def case( self, keras_model: tf.keras.Model, qspec: QuantizationSpec ) - QuantizationSpec: raise NotImplementedError(case method must be overridden by user)第二步重写两个方法1.case必选—— 固定签名如下工具包会自动为quantize_model函数custom_qdq_cases参数中的每个成员调用case方法(function) CustomQDQInsertionCase.case( self, keras_model : tf.keras.Model, qspec : QuantizationSpec ) - QuantizationSpec修改默认层行为的逻辑都应写在该函数中并且必须返回一个QuantizationSpec对象。从 quantize.py 的实现可以看到工具包对返回值有硬性校验qspec_case_object custom_qdq_case.case(model, quantization_spec) if qspec_case_object: assert isinstance( qspec_case_object, QuantizationSpec ), [E] {} does not return an object of QuantizationSpec..format(...)2.info可选—— 辅助方法用字符串说明case方法内部的逻辑便于阅读与调试。基类的默认实现返回空字符串。第三步注册到量化流程实例化新类并加入列表传给quantize_model的custom_qdq_cases参数q_nn_model quantize_model( modelnn_model_original, custom_qdq_cases[MyCustomQDQCase()] )完整示例MaxPoolQDQCase以下是文档提供的完整示例与 custom_qdq_cases.py 中的实现一致class MaxPoolQDQCase(CustomQDQInsertionCase): def __init__(self) - None: super().__init__() def info(self) - str: return Enables quantization of MaxPool layers. def case( self, keras_model: tf.keras.Model, qspec: QuantizationSpec ) - QuantizationSpec: mp_qspec QuantizationSpec() for layer in keras_model.layers: if isinstance(layer, tf.keras.layers.MaxPooling2D): if check_is_quantizable_by_layer_name(qspec, layer.name): mp_qspec.add( namelayer.name, quantize_inputTrue, quantize_weightFalse ) return mp_qspec对该示例的逐行解读case方法遍历keras_model.layers用isinstance筛选出所有MaxPooling2D层对每个 MaxPool 层调用check_is_quantizable_by_layer_name(qspec, layer.name)检查该层是否没有出现在用户传入的QuantizationSpec中通过mp_qspec.add(...)为符合条件的层追加LayerConfigquantize_inputTrue, quantize_weightFalse因为池化层无权重最后返回新的mp_qspec。注意这里构造的是一个全新的QuantizationSpec而不是修改传入的qspec。check_is_quantizable_by_layer_name是实现用户规格优先的关键辅助函数其语义是如果current_layer_name不在用户传入的qspec中则返回True可由自动策略接管否则返回False尊重用户的显式配置。实现见 custom_qdq_cases.py。关于qspec参数的使用契约case方法中的逻辑可以使用、也可以不使用用户传入的QuantizationSpec对象但它必须返回一个新的QuantizationSpec用以承载更新后的层行为。以上述MaxPoolQDQCase为例其逻辑就依赖用户传入的QuantizationSpeccheck_is_quantizable_by_layer_name检查层名是否在用户对象中并优先遵循用户的规格用户已显式配置的层自动策略不再覆盖。编写更复杂策略时的辅助工具在 custom_qdq_cases.py 中还提供了两个判断父层模式的工具函数可用于编写依赖网络拓扑的复杂 Caseis_parent_type(parent_class: str, class_typeConv)判断父层类名是否属于某类型如Conv、BatchNorm、Dropout、Activationis_parent_pattern(parent_info: dict, pattern: List)判断父层谱系是否匹配给定模式例如[BatchNorm, Conv]表示父层是 BatchNorm、祖父层是 Conv对应 ResNet-v1 中Conv2D - BatchNormalization - Add的间接连接utils.find_my_predecessors(keras_model, layer_name)获取某层的全部前驱层信息ResidualConnectionQDQCase正是借助它遍历Add层的各输入分支找出带 Conv 模式且只有一个输出分支的路径并跳过量化见 custom_qdq_cases.py。最佳实践与注意事项必须返回QuantizationSpec只要编写了CustomQDQInsertionCase其case方法就必须返回QuantizationSpec对象否则quantize_model会触发断言错误quantize.py。用户规格优先在case中应通过check_is_quantizable_by_layer_name检查层名是否已被用户显式配置避免自动策略覆盖用户的明确意图。默认返回新对象而非修改入参case应当构造并返回一个新的QuantizationSpecquantize_model内部会把所有 Case 的输出与用户传入的quantization_spec合并layers.extend。模型类型限制工具包目前仅支持 KerasFunctional或Sequential模型Subclassed模型会在quantize_model入口被断言拒绝quantize.py。量化模式的关系Custom Q/DQ Insertion Case常与Full Custom Quantization模式配合也可用于 Partial 模式。关于Full Default / Full Custom / Partial三种模式的完整定义参见 basics.md。优先复用内置 Case编写新 Case 前先检查 custom_qdq_cases.py 中已有的通用 CaseResidualConnectionQDQCase、BNQDQCase、MaxPoolQDQCase并通过组合方式构建网络特定策略如ResNetV2QDQCase的做法。用测试验证行为可以参考 tests/custom_qdq_cases_test.py 的断言模式检查自定义 Case 返回的LayerConfig.name与quantization_index是否符合预期。延伸阅读自定义 Q/DQ 插入的完整官方说明add_custom_qdq_cases.md量化模式Full/Partial与术语定义basics.md受支持层及默认量化行为矩阵add_new_layer_support.md内置 Case 完整实现custom_qdq_cases.py基类CustomQDQInsertionCase定义custom_qdq_case_base.pyquantize_model与QuantizationSpec实现quantize.py内置 Case 的单元测试custom_qdq_cases_test.py各模型 QAT 精度与延迟结果model_zoo.md【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考