convert_fx

classtorch.ao.quantization.quantize_fx.convert_fx(graph_module, convert_custom_config=None, _remove_qconfig=True, qconfig_mapping=None, backend_config=None)[源代码]

将校准或训练后的模型转换为量化模型

参数
  • graph_module (*) – 一个已准备好并经过校准/训练的模型(GraphModule)

  • convert_custom_config (*) – 用于 convert 函数的自定义配置。更多详情请参见 ConvertCustomConfig

  • _remove_qconfig (*) – 在转换完成后,选择是否从模型中移除 qconfig 属性。

  • qconfig_mapping (*) –

    用于指定将模型转换为量化模型的配置。

    键必须包含在传递给 prepare_fxprepare_qat_fxqconfig_mapping 中定义的所有键,并具有相同的值或 None。还可以指定额外的键,其值设置为 None

    对于值为None的每个条目,我们在模型中跳过对其进行量化:

    qconfig_mapping = QConfigMapping
        .set_global(qconfig_from_prepare)
        .set_object_type(torch.nn.functional.add, None)  # skip quantizing torch.nn.functional.add
        .set_object_type(torch.nn.functional.linear, qconfig_from_prepare)
        .set_module_name("foo.bar", None)  # skip quantizing module "foo.bar"
    
    • backend_config (BackendConfig): 后端的配置,描述了相关设置

      应在后端对操作符进行量化,这包括支持不同的量化模式(如静态、动态和仅权重)、数据类型(如quint8和qint8)以及每个操作符的观察器放置。此外,还需要考虑融合操作符的情况。详情请参见BackendConfig

返回值

一个量化的 torch.nn.Module 模型

返回类型

GraphModule

示例:

# prepared_model: the model after prepare_fx/prepare_qat_fx and calibration/training
# convert_fx converts a calibrated/trained model to a quantized model for the
# target hardware, this includes converting the model first to a reference
# quantized model, and then lower the reference quantized model to a backend
# Currently, the supported backends are fbgemm (onednn), qnnpack (xnnpack) and
# they share the same set of quantized operators, so we are using the same
# lowering procedure
#
# backend_config defines the corresponding reference quantized module for
# the weighted modules in the model, e.g. nn.Linear
# TODO: add backend_config after we split the backend_config for fbgemm and qnnpack
# e.g. backend_config = get_default_backend_config("fbgemm")
quantized_model = convert_fx(prepared_model)
本页目录