torch.ao.ns._numeric_suite_fx

警告

此模块是一个早期原型，可能会发生变化。

此模块包含用于比较模型权重和激活的工具。示例用法：

import copy
import torch
import torch.ao.quantization.quantize_fx as quantize_fx
import torch.ao.ns._numeric_suite_fx as ns

m = torch.nn.Sequential(torch.nn.Conv2d(1, 1, 1)).eval()
mp = quantize_fx.prepare_fx(m, {'': torch.ao.quantization.default_qconfig})
# We convert a copy because we need the original prepared model
# to be available for comparisons, and `quantize_fx.convert_fx` is inplace.
mq = quantize_fx.convert_fx(copy.deepcopy(mp))

#
# Comparing weights
#

# extract weight pairs
weight_comparison = ns.extract_weights('a', mp, 'b', mq)

# add SQNR for each comparison, inplace
ns.extend_logger_results_with_comparison(
    weight_comparison, 'a', 'b', torch.ao.ns.fx.utils.compute_sqnr,
    'sqnr')

# weight_comparison contains the weights from `mp` and `mq` stored
# in pairs, and can be used for further analysis.


#
# Comparing activations, with error propagation
#

# add loggers
mp_ns, mq_ns = ns.add_loggers(
    'a', copy.deepcopy(mp),
    'b', copy.deepcopy(mq),
    ns.OutputLogger)

# send an example datum to capture intermediate activations
datum = torch.randn(1, 1, 1, 1)
mp_ns(datum)
mq_ns(datum)

# extract intermediate activations
act_comparison = ns.extract_logger_info(
    mp_ns, mq_ns, ns.OutputLogger, 'b')

# add SQNR for each comparison, inplace
ns.extend_logger_results_with_comparison(
    act_comparison, 'a', 'b', torch.ao.ns.fx.utils.compute_sqnr,
    'sqnr')

# act_comparison contains the activations from `mp_ns` and `mq_ns` stored
# in pairs, and can be used for further analysis.

#
# Comparing activations, without error propagation
#

# create shadow model
mp_shadows_mq = ns.add_shadow_loggers(
    'a', copy.deepcopy(mp),
    'b', copy.deepcopy(mq),
    ns.OutputLogger)

# send an example datum to capture intermediate activations
datum = torch.randn(1, 1, 1, 1)
mp_shadows_mq(datum)

# extract intermediate activations
shadow_act_comparison = ns.extract_shadow_logger_info(
    mp_shadows_mq, ns.OutputLogger, 'b')

# add SQNR for each comparison, inplace
ns.extend_logger_results_with_comparison(
    shadow_act_comparison, 'a', 'b', torch.ao.ns.fx.utils.compute_sqnr,
    'sqnr')

# shadow_act_comparison contains the activations from `mp_ns` and `mq_ns` stored
# in pairs, and can be used for further analysis.

class torch.ao.ns._numeric_suite_fx.OutputLogger(ref_node_name, prev_node_name, model_name, ref_name, prev_node_target_type, ref_node_target_type, results_type, index_within_arg, index_of_arg, fqn, qconfig_str='')

捕获中间值的基类。

forward(x)

class torch.ao.ns._numeric_suite_fx.OutputComparisonLogger(*args, **kwargs)

与 OutputLogger 相同，但还需要原始激活信息，以便在校准时间计算比较

forward(x, x_ref)

class torch.ao.ns._numeric_suite_fx.NSTracer(skipped_module_names, skipped_module_classes)

就像常规的 FX 量化跟踪器一样，但将观察者和 fake_quantize 模块视为叶子模块。

is_leaf_module(m, module_qualified_name)

返回类型:

布尔型

torch.ao.ns._numeric_suite_fx.extract_weights(model_name_a, model_a, model_name_b, model_b, base_name_to_sets_of_related_ops=None, unmatchable_types_map=None, op_to_type_to_weight_extraction_fn=None)

从模型 A 和模型 B 中提取权重，并返回比较结果。

参数:

• model_name_a (str) – 用于结果中的模型 A 的字符串名称

• 模型_a（模块）- 模型 A

• 模型名_b（str）- 用于结果中的模型 B 的字符串名称

• 模型_b（模块）- 模型 B

• base_name_to_sets_of_related_ops（Optional[dict[str, set[Union[Callable, str]]]]）- 可选的子图基本节点覆盖，可能发生变化

• unmatchable_types_map (Optional[dict[str, set[Union[Callable, str]]]]) – 可选的未匹配类型覆盖，可能发生变化

• op_to_type_to_weight_extraction_fn (Optional[dict[str, dict[Callable, Callable]]]) – 可选的从类型提取权重的函数覆盖，可能发生变化

返回值:

NSResultsType，包含权重比较

返回类型:

dict[str, dict[str, dict[str, list[dict[str, Any]]]]]

torch.ao.ns._numeric_suite_fx.add_loggers(name_a, model_a, name_b, model_b, logger_cls, should_log_inputs=False, base_name_to_sets_of_related_ops=None, unmatchable_types_map=None)

使用日志记录器对模型 A 和模型 B 进行配置。

参数:

• name_a (str) – 要在结果中使用模型 A 的字符串名称

• model_a (Module) – 模型 A

• name_b (str) – 要在结果中使用的模型 B 的字符串名称

• model_b (Module) – 模型 B

• logger_cls (Callable) – 要使用的 Logger 类

• base_name_to_sets_of_related_ops (Optional[dict[str, set[Union[Callable, str]]]]) – 可选的子图基本节点覆盖，可能发生变化

• unmatchable_types_map (Optional[dict[str, set[Union[Callable, str]]]]) – 可选的未匹配类型覆盖，可能发生变化

返回值:

返回一个包含 (model_a_with_loggers, model_b_with_loggers) 的元组。就地修改两个模型。

返回类型:

tuple[torch.nn.modules.module.Module, torch.nn.modules.module.Module]

torch.ao.ns._numeric_suite_fx.extract_logger_info(model_a, model_b, logger_cls, model_name_to_use_for_layer_names)

遍历 model_a 和 model_b 中的所有记录器，并提取记录的信息。

参数:

• model_a（模块）- 模型 A

• model_b（模块）- 模型 B

• logger_cls（可调用）- 要使用的 Logger 类

• 要使用的模型名称（str）- 输出中用于层名称的字符串名称

返回值:

包含已记录比较的 NSResultsType

返回类型:

dict[str, dict[str, dict[str, list[dict[str, Any]]]]]

torch.ao.ns._numeric_suite_fx.add_shadow_loggers(name_a, model_a, name_b, model_b, logger_cls, should_log_inputs=False, base_name_to_sets_of_related_ops=None, node_type_to_io_type_map=None, unmatchable_types_map=None)

使用阴影记录器对 A 型号和 B 型号进行仪器建模。

参数:

• name_a (str) – 用于结果中的模型 A 的字符串名称

• model_a (模块) – 模型 A

• name_b (str) – 用于结果中的模型 B 的字符串名称

• 模型_b（模块）- 模型 B

• logger_cls（可调用）- 要使用的 Logger 类

• should_log_inputs（布尔值）- 是否记录输入

• base_name_to_sets_of_related_ops（Optional[dict[str, set[Union[Callable, str]]]]) – 可选的子图基本节点覆盖，可能发生变化

• unmatchable_types_map (Optional[dict[str, set[Union[Callable, str]]]]) – 可选的未匹配类型覆盖，可能发生变化

返回类型:

模块

torch.ao.ns._numeric_suite_fx.extract_shadow_logger_info(model_a_shadows_b, logger_cls, model_name_to_use_for_layer_names)

遍历阴影模型中的所有日志记录器，并提取记录的信息。

参数:

• model_a_shadows_b (Module) – 阴影模型

• logger_cls (Callable) – Logger 类

• model_name_to_use_for_layer_names (str) – 用于输出层名称的模型字符串名称

返回值:

NSResultsType，包含已记录的比较

返回类型:

dict[str, dict[str, dict[str, list[dict[str, Any]]]]]

torch.ao.ns._numeric_suite_fx.extend_logger_results_with_comparison(results, model_name_1, model_name_2, comparison_fn, comparison_name)

将模型_name_2 的日志值与模型_name_1 的对应值进行比较，使用 comparison_fn。将结果记录在 model_name_2 的结果中，名称为 comparison_name。就地修改 results。

参数:

• results (dict[str, dict[str, dict[str, list[dict[str, Any]]]]]) – 从 extract_logger_info 或 extract_shadow_logger_info 获取的结果数据结构。

• model_name_1 (str) – 模型 1 的字符串名称

• model_name_2 (str) – 模型 2 的字符串名称

• comparison_fn (Callable[[Tensor, Tensor], Tensor]) – 比较两个张量的函数

• comparison_name (str) – 用于输出层名称的模型的字符串名称

torch.ao.ns._numeric_suite_fx.prepare_n_shadows_model(model, example_inputs, qconfig_multi_mapping, backend_config, custom_prepare_fn=None, custom_prepare_kwargs=None, custom_tracer=None)

给定一个包含 M 个操作的图模型，

args_kwargs_m -> op_m -> output_m

并且为每个操作设置 N 个 qconfig，创建一个新的模型，其中每个 op_m 的子图都被转换成

     |---------> op_m_n -> log_m_n
     |                     /
args_kwargs_m ---------> op_m -> log_m_0

其中 op_m_n 是包含 qconfig_n 转换的 op_m 的子模块，其内部图看起来像

args_m -------- op_m_prepared_with_qconfig_n -> out_m_n
            /
kwargs_m ---

这对于在单次模型遍历中测试多个层的不同量化非常有用。

未来 PR 的高级别待办事项：* 想出一个更好的方法来命名输出结构 * 返回结果数据结构而不是打印它 * 在文档块中添加示例

返回类型:

GraphModule

torch.ao.ns._numeric_suite_fx.loggers_set_enabled(model, enabled)

设置模型日志记录器的启用设置

torch.ao.ns._numeric_suite_fx.loggers_set_save_activations(model, save_activations)

设置模型日志记录器的 save_activations 设置

torch.ao.ns._numeric_suite_fx.convert_n_shadows_model(model, custom_convert_fn=None, custom_convert_kwargs=None)

给定从 prepare_n_shadows_model 准备好的模型，对每个阴影子模块运行 convert_fx

返回类型:

GraphModule

torch.ao.ns._numeric_suite_fx.extract_results_n_shadows_model(model)

从模型中提取日志结果。

返回类型:

dict[str, dict[str, dict[str, list[dict[str, Any]]]]]

torch.ao.ns._numeric_suite_fx.print_comparisons_n_shadows_model(results)

打印提取结果的摘要。

torch.ao.ns.fx.utils

警告

此模块是一个早期原型，可能会发生变化。

torch.ao.ns.fx.utils.compute_sqnr(x, y)

torch.ao.ns.fx.utils.compute_normalized_l2_error(x, y)

torch.ao.ns.fx.utils.compute_cosine_similarity(x, y)