modelopt.onnx.quantization.quantize

quantize(onnx_path, quantize_mode='int8', calibration_data=None, calibration_method=None, calibration_cache_path=None, calibration_shapes=None, calibration_eps=['cuda:0', 'cpu', 'trt'], op_types_to_quantize=None, op_types_to_exclude=None, nodes_to_quantize=None, nodes_to_exclude=None, use_external_data_format=False, keep_intermediate_files=False, output_path=None, verbose=False, trt_plugins=None, trt_plugins_precision=None, high_precision_dtype=None, mha_accumulation_dtype='fp32', disable_mha_qdq=False, dq_only=True, block_size=None, use_zero_point=False, **kwargs)

量化提供的ONNX模型。

Parameters:

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

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

  • calibration_data (ndarray | Dict[str, ndarray]) – 校准数据，可以是numpy数组或numpy数组的列表/字典。

  • calibration_method (str) – 校准方法选择。选项有 int8: ‘entropy’（默认）和 ‘max’, fp8: ‘max’（默认）和 int4: ‘awq_clip’（默认）, ‘awq_lite’, ‘awq_full’ 和 ‘rtn_dq’.

  • calibration_cache_path (str) – 预计算的激活张量范围的路径，也称为校准缓存。

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

  • op_types_to_quantize (List[str]) – 要量化的操作类型列表。如果为None（默认值），则所有支持的操作符都会被量化。 此标志不支持正则表达式。

  • op_types_to_exclude (List[str]) – 要从量化中排除的操作类型列表。此标志不支持正则表达式。

  • nodes_to_quantize (List[str]) – 要量化的节点名称列表。如果为None（默认值），则所有支持的节点都会被量化。 此标志支持正则表达式。

  • nodes_to_exclude (List[str]) – 要从量化中排除的节点名称列表。此标志支持正则表达式。

  • use_external_data_format (bool) – 如果为True，将使用单独的数据路径来存储量化模型的权重。

  • keep_intermediate_files (bool) – 如果为True，保留在ONNX模型转换/校准过程中生成的所有中间文件。

  • output_path (str) – 输出文件名以保存量化后的ONNX模型。 如果为None，则保存在与原始ONNX模型相同的目录中，并添加.quant后缀。

  • verbose (bool) – 如果为True，则在量化过程中打印节点分区、选择等的详细信息。

  • trt_plugins (str) – 指定自定义TensorRT插件库的路径，格式为.so（编译后的共享库）。 对于多个路径，用分号分隔，例如：“lib_1.so;lib_2.so”。 如果此项不为None或模型包含自定义操作，TensorrtExecutionProvider将作为校准执行提供者的首选，这意味着TensorRT是必需的。

  • trt_plugins_precision (List[str]) – 一个以空格分隔的列表，指示每个自定义操作的精度。 每个项目应具有格式 :，其中精度可以是 fp32（默认）或 fp16。 例如：op_type_1:fp16 op_type_2:fp32。

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

  • mha_accumulation_dtype (str) – MHA 累加数据类型。可选值为 [‘fp32’, ‘fp16’]。默认为 ‘fp32’。 如果 quantize_mode == ‘fp8’ 且 high_precision_dtype == ‘fp32’，则会在 MHA 的 bmm1 和 bmm2 的输入和输出张量上添加 Cast 节点。

  • disable_mha_qdq (bool) – 不要在MHA模式中的MatMuls上添加Q/DQ层。

  • dq_only (bool) – 如果为True（默认），则仅向模型添加DQ节点。如果为False，则向模型添加Q/DQ节点。

  • block_size (int | None) – int4量化的块大小参数。

  • use_zero_point (bool) – 如果为True，则使用基于零点的量化。

  • kwargs (Any) – 用于int4量化的额外关键字参数，包括： - awqlite_alpha_step (float): lite的alpha步长，范围[0, 1]。 - awqclip_alpha_step (float): clip的最小alpha步长，范围[awqclip_alpha_step, 1]。 - awqclip_alpha_min (float): 用于找到clip最佳alpha的alpha步长。 - awqclip_bsz_col (int): 处理clip中列维度的批量大小。

  • calibration_shapes (str) –

Returns:

无，将量化后的onnx模型写入提供的output_path中，或者写入相同目录，文件名为“.quant.onnx”。

Return type:

无