graph_utils

提供用于QDQ放置的ONNX图相关工具。

函数

add_fp16_fp32_cast	将cast_to_fp16节点添加到层的输入，并将cast_to_fp32添加到输出。
build_non_residual_input_map	构建一个非残差Add输入名称到给定图中的Add节点名称的映射。
classify_partition_nodes	我们应该对分区节点进行部分量化，这些节点的输入位于分区之外。
convert_fp16_io	将图形I/O转换为FP16。
expand_node_names_from_patterns	从给定的模式中扩展节点名称。
filter_quantizable_kgen_heads	如果遵循CASK分区，则返回kgen头名称列表。
find_fp8_mha_partitions	匹配FP8 MHA: Q -> DQ -> BMM1 -> (Mul/Div) -> (Add) -> Softmax -> (Cast) -> Q -> DQ -> BMM2 -> Q -> DQ.
find_mha_partitions	匹配 MHA: BMM1 -> (乘法/除法) -> (加法) -> Softmax -> (类型转换) -> BMM2。
find_nodes_from_matmul_to_exclude	找到符合gemv条件以排除的MatMul节点。
find_nodes_from_mha_to_exclude	在MHA模式中查找MatMul节点以排除。
find_nodes_to_exclude	从ONNX图中找到与用户排除模式匹配的节点名称。
get_fusible_backbone	如果给定节点匹配模式，则返回其线性骨干节点。
get_tensor_consumer_nodes	返回一个张量名称及其消费者节点对象映射的字典。
get_tensor_producer_nodes	返回一个张量名称及其生产者节点对象映射的字典。
has_const_input	返回给定节点是否有任何常量输入。
has_path_type	检查给定节点是否是给定前向/后向路径类型的开始/结束。
insert_fp8_mha_casts	插入三个类型转换操作。
insert_matmul_casts	为MatMul的两个输入和输出插入三个转换节点。
is_const_input	返回给定的张量是否是一个初始化器或由可常量折叠节点产生。
print_stat	收集并打印量化模型的统计信息。
remove_partial_input_qdq	通过从标记的输入中移除QDQ节点来修改onnx模型，例如。

add_fp16_fp32_cast(onnx_path, custom_ops_to_cast_to_fp16)

将cast_to_fp16节点添加到层的输入，并将cast_to_fp32添加到输出。

build_non_residual_input_map(graph)

构建一个非残差Add输入名称到给定图中的Add节点名称的映射。

这假设Add层只有2个输入。

我们将把具有卷积节点且其单个输出与另一个非常量输入张量（逐元素）相加的子图称为“残差加法”子图，因为它出现在使用残差连接的现代卷积网络中。

Parameters:

graph (Graph) – Onnx 模型图。

Returns:

添加节点名称与其非残差输入名称的字典。

Return type:

Dict[str, str]

classify_partition_nodes(partitions)

我们应该对分区节点进行部分量化，这些节点的输入位于分区之外。

Parameters:

partitions (List[List[Node]]) – 由modelopt ptq算法创建的分区。

Returns:

不可量化节点列表。 可量化节点列表。 部分可量化输入列表，包含不可量化输入信息（源节点，目标节点，输入名称）

Return type:

元组[列表[节点], 列表[节点], 列表[元组[节点, 节点, 字符串]]]

convert_fp16_io(graph)

将图形I/O转换为FP16。

expand_node_names_from_patterns(graph, name_patterns)

从给定的模式中扩展节点名称。

Parameters:

• 图 (GraphProto | Graph) –

• name_patterns (List[str]) –

Return type:

列表[字符串]

filter_quantizable_kgen_heads(cask_fusible_partitions, kgen_partitions, quantizable_op_types)

如果遵循CASK分区，则返回kgen头名称列表。

Parameters:

• cask_fusible_partitions (List[List[Node]]) –

• kgen_partitions (List[List[Node]]) –

• quantizable_op_types (List[str]) –

Return type:

元组[列表[节点], 列表[元组[节点, 节点, 字符串]]]

find_fp8_mha_partitions(graph)

匹配FP8 MHA: Q -> DQ -> BMM1 -> (Mul/Div) -> (Add) -> Softmax -> (Cast) -> Q -> DQ -> BMM2 -> Q -> DQ.

find_mha_partitions(graph)

匹配 MHA: BMM1 -> (乘法/除法) -> (加法) -> Softmax -> (类型转换) -> BMM2。

find_nodes_from_matmul_to_exclude(onnx_path, use_external_data_format=False, intermediate_generated_files=None, calibration_shapes=None, calibration_eps=['cuda:0', 'cpu', 'trt'], verbose=False)

找到符合gemv条件以排除的MatMul节点。

在matmul中，m或n中的任何一个为1时，此matmul无法利用TensorCores。在TRT中添加Q/DQ层的性能不佳。因此，在这种情况下，不要向此matmul添加Q/DQ层。

Parameters:

• onnx_path (str) – ONNX模型的路径。

• use_external_data_format (bool) – 如果不为None，此路径将用于存储量化模型的权重。

• intermediate_generated_files (List[str]) – 量化后将删除的中间生成文件列表。

• calibration_shapes (str) – 用于推理的模型输入形状。

• calibration_eps (List[str]) – 用于校准模型的执行提供程序（EP）的优先级顺序。 可以是 [‘cuda:x’, ‘cpu’, ‘trt’] 的任意子集，其中 ‘x’ 是设备ID。

• verbose (bool) – 如果为True，则打印要排除的matmul节点。

Returns:

要从量化中排除的节点列表。

Return type:

列表[字符串]

find_nodes_from_mha_to_exclude(onnx_path, use_external_data_format=False, nodes_to_exclude=None, disable_mha_qdq=False, quantize_mode='int8', high_precision_dtype=None, intermediate_generated_files=None, calibration_shapes=None, calibration_eps=['cuda:0', 'cpu', 'trt'], verbose=False)

在MHA模式中查找MatMul节点以排除。

如果设置了disable_mha_qdq，则不要在MHA模式中的MatMuls上添加Q/DQ层。 否则，当quantize_mode == “fp8”且high_precision_dtype == “fp16”时， 如果head_size不是16的倍数或MHA模式中存在maskadd， 则不要在MHA模式中的MatMuls上添加Q/DQ层。 否则，当quantize_mode == “int8”时，如果seq_len > 512，则不要在MHA模式中的MatMuls上添加Q/DQ层。

Parameters:

• onnx_path (str) – ONNX模型的路径。

• use_external_data_format (bool) – 如果不为None，此路径将用于存储量化模型的权重。

• nodes_to_exclude (List[str]) – 要从量化中排除的节点列表。

• disable_mha_qdq (bool) – 如果为True，所有MHA的BMM1和BMM2将被添加到nodes_to_exclude中。 否则，当is_fp8fp16为True时，将检查每个MHA是否启用QDQ。

• quantize_mode (str) – 量化模式。可选值为 'int8'（默认）、'int4' 和 'fp8'。

• high_precision_dtype (str) – 高精度数据类型，可选值为 [‘fp32’, ‘fp16’]。如果 high_precision_dtype == ‘fp16’，模型的权重和激活值将被转换为 fp16。

• intermediate_generated_files (List[str]) – 量化后将删除的中间生成文件列表。

• calibration_shapes (str) – 用于推理的模型输入形状。

• calibration_eps (List[str]) – 用于校准的执行提供者（EP）的优先级列表。

• verbose (bool) – 如果为True，打印要排除的matmul节点。

Returns:

要从量化中排除的节点列表。

Return type:

列表[字符串]

find_nodes_to_exclude(graph, nodes_to_exclude, op_types_to_exclude)

从ONNX图中找到与用户排除模式匹配的节点名称。

Parameters:

• 图 (图) –

• nodes_to_exclude (List[str]) –

• op_types_to_exclude (List[str]) –

get_fusible_backbone(node, graph)

如果给定节点匹配模式，则返回其线性骨干节点。

TensorRT 在某些特定模式下会将卷积与BN、Relu等融合。此规则尝试匹配其中一些模式。注意：BiasAdd 和 ConstMul 在路径类型中是可选的。

Parameters:

• node (Node) – 模式的起始节点。

• graph (Graph) – ONNX 模型图。

Returns:

给定节点的骨干节点，如果未找到则为None。

Return type:

节点 | 无

get_tensor_consumer_nodes(graph)

返回一个张量名称及其消费者节点对象映射的字典。

Parameters:

graph (GraphProto) – ONNX 模型图。

Returns:

字典，键是张量名称，值是它们的消费者节点对象

Return type:

字典[字符串, 列表[节点原型]]

get_tensor_producer_nodes(graph)

返回一个张量名称及其生产者节点对象映射的字典。

注意。我们创建了一个特殊的根类型节点作为外部输入生产者，以便于实现。

Parameters:

graph (GraphProto) – ONNX 模型图。

Returns:

字典，键是张量名称，值是它们的生产者节点对象

Return type:

Dict[str, NodeProto]

has_const_input(node)

返回给定节点是否有任何常量输入。

Parameters:

节点 (Node) –

Return type:

bool

has_path_type(node, graph, path_type, is_forward, wild_card_types=[], path_nodes=[])

检查给定节点是否是给定前向/后向路径类型的开始/结束。

注意，路径可以相对于节点向前或向后，取决于下一级节点。 此外，此方法可以与可选节点一起使用，并收集遍历的路径。

Parameters:

• 节点 (Node) – 路径的起始节点。

• graph (Graph) – ONNX 模型图。

• path_type (List[str]) – 从给定节点匹配的路径类型。

• is_forward (bool) – 是否匹配正向或反向路径。

• wild_card_types (List[str]) – 通配符类型，这些类型的节点会被跳过，不会与路径类型匹配。

• path_nodes (List[Node]) – 匹配路径中累积的节点。

Returns:

布尔值，表示给定节点是否为给定前向/后向路径类型的起点/终点。

Return type:

bool

insert_fp8_mha_casts(onnx_model)

插入三个类型转换操作。

第一个转换将在MatMul的input0之前添加，以将fp16转换为fp32。 第二个转换将在MatMul的input1之前添加，以将fp16转换为fp32。 第三个转换将在MatMul的输出之后添加，以将fp32转换回fp16。 在FP8 MHA部分插入Cast操作实际上禁止了MHA以FP16累加运行，因为编译器仅为FP8 MHA提供了FP32累加内核。

insert_matmul_casts(graph, matmul_node)

为MatMul的两个输入和输出插入三个转换节点。

is_const_input(tensor)

返回给定的张量是否是一个初始化器或由可常量折叠节点产生。

Parameters:

张量 (Tensor) –

Return type:

bool

print_stat(graph, verbose)

收集并打印量化模型的统计信息。

Parameters:

• 图 (图) –

• verbose (bool) –

Return type:

无

remove_partial_input_qdq(graph, no_quantize_inputs)

通过从标记的输入中移除QDQ节点来修改onnx模型，例如非残差输入等。

Parameters:

• graph (Graph) – Onnx 模型图。

• no_quantize_inputs (List[Tuple[Node, Node, str]]) – 列出不可量化的输入信息，格式为 (src, dst, input_name)

Return type:

无