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})
# 我们转换一个副本，因为我们需要的原始准备模型
# 用于比较，而`quantize_fx.convert_fx`是就地操作。
mq = quantize_fx.convert_fx(copy.deepcopy(mp))

#
# 比较权重
#

# 提取权重对
weight_comparison = ns.extract_weights('a', mp, 'b', mq)

# 为每个比较添加SQNR，就地操作
ns.extend_logger_results_with_comparison(
    weight_comparison, 'a', 'b', torch.ao.ns.fx.utils.compute_sqnr,
    'sqnr')

# weight_comparison包含来自`mp`和`mq`的权重
# 存储为对，可用于进一步分析。


#
# 比较激活值，带误差传播
#

# 添加记录器
mp_ns, mq_ns = ns.add_loggers(
    'a', copy.deepcopy(mp),
    'b', copy.deepcopy(mq),
    ns.OutputLogger)

# 发送一个示例数据以捕获中间激活值
datum = torch.randn(1, 1, 1, 1)
mp_ns(datum)
mq_ns(datum)

# 提取中间激活值
act_comparison = ns.extract_logger_info(
    mp_ns, mq_ns, ns.OutputLogger, 'b')

# 为每个比较添加SQNR，就地操作
ns.extend_logger_results_with_comparison(
    act_comparison, 'a', 'b', torch.ao.ns.fx.utils.compute_sqnr,
    'sqnr')

# act_comparison包含来自`mp_ns`和`mq_ns`的激活值
# 存储为对，可用于进一步分析。

#
# 比较激活值，不带误差传播
#

# 创建影子模型
mp_shadows_mq = ns.add_shadow_loggers(
    'a', copy.deepcopy(mp),
    'b', copy.deepcopy(mq),
    ns.OutputLogger)

# 发送一个示例数据以捕获中间激活值
datum = torch.randn(1, 1, 1, 1)
mp_shadows_mq(datum)

# 提取中间激活值
shadow_act_comparison = ns.extract_shadow_logger_info(
    mp_shadows_mq, ns.OutputLogger, 'b')

# 为每个比较添加SQNR，就地操作
ns.extend_logger_results_with_comparison(
    shadow_act_comparison, 'a', 'b', torch.ao.ns.fx.utils.compute_sqnr,
    'sqnr')

# shadow_act_comparison包含来自`mp_ns`和`mq_ns`的激活值
# 存储为对，可用于进一步分析。

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)[源代码]¶

Return type: bool

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中提取权重，并返回一个比较结果。

Parameters

model_name_a (str) – 在结果中使用的模型A的字符串名称
model_a (模块) – 模型 A
model_name_b (str) – 在结果中使用的模型B的字符串名称
model_b (模块) – 模型 B
base_name_to_sets_of_related_ops (可选[字典[字符串, 集合[联合[可调用, 字符串]]]]) – 可选的子图基础节点覆盖，可能会更改
unmatchable_types_map (可选[字典[字符串, 集合[联合[可调用, 字符串]]]]) – 不可匹配类型的可选覆盖，可能会更改
op_to_type_to_weight_extraction_fn (可选[字典[字符串, 字典[可调用, 可调用]]]) – 从类型中提取权重的函数的可选覆盖，可能会更改

Returns

NSResultsType，包含重量比较

Return type

字典[字符串, 字典[字符串, 字典[字符串, 列表[字典[字符串, 任意]]]]]

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。

Parameters

name_a (str) – 在结果中使用的模型A的字符串名称
model_a (模块) – 模型 A
name_b (str) – 结果中使用的模型B的字符串名称
model_b (模块) – 模型 B
logger_cls (可调用) – 要使用的日志记录器类
base_name_to_sets_of_related_ops (可选[字典[字符串, 集合[联合[可调用, 字符串]]]]) – 可选的子图基础节点覆盖，可能会更改
unmatchable_types_map (可选[字典[字符串, 集合[联合[可调用, 字符串]]]]) – 不可匹配类型的可选覆盖，可能会更改

Returns

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

Return type

元组[模块, 模块]

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 中的所有日志记录器，并提取记录的信息。

Parameters

model_a (模块) – 模型 A
model_b (模块) – 模型 B
logger_cls (可调用) – 要使用的日志记录器类
model_name_to_use_for_layer_names (str) – 用于输出层名称的模型字符串名称

Returns

NSResultsType，包含记录的比较结果

Return type

字典[字符串, 字典[字符串, 字典[字符串, 列表[字典[字符串, 任意]]]]]

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。

Parameters

name_a (str) – 在结果中使用的模型A的字符串名称
model_a (模块) – 模型 A
name_b (str) – 结果中使用的模型B的字符串名称
model_b (模块) – 模型 B
logger_cls (可调用对象) – 要使用的日志记录器类
should_log_inputs (bool) – 是否记录输入
base_name_to_sets_of_related_ops (可选[字典[字符串, 集合[联合[可调用, 字符串]]]]) – 可选的子图基础节点覆盖，可能会更改
unmatchable_types_map (可选[字典[字符串, 集合[联合[可调用, 字符串]]]]) – 不可匹配类型的可选覆盖，可能会更改

Return type

模块

torch.ao.ns._numeric_suite_fx.extract_shadow_logger_info(model_a_shadows_b, logger_cls, model_name_to_use_for_layer_names)[源代码]¶

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

Parameters

model_a_shadows_b (模块) – 影子模型
logger_cls (可调用) – 要使用的日志记录器类
model_name_to_use_for_layer_names (str) – 用于输出层名称的模型字符串名称

Returns

NSResultsType，包含记录的比较结果

Return type

字典[字符串, 字典[字符串, 字典[字符串, 列表[字典[字符串, 任意]]]]]

torch.ao.ns._numeric_suite_fx.extend_logger_results_with_comparison(results, model_name_1, model_name_2, comparison_fn, comparison_name)[源代码]¶

将model_name_2中的记录值与model_name_1中的相应值进行比较，使用comparison_fn。将结果记录在model_name_2的结果中，位于comparison_name下。就地修改results。

Parameters

结果 (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 (可调用[[张量, 张量], 张量]) – 用于比较两个张量的函数
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 是 op_m 封装在子模块中并通过 qconfig_n 进行转换，其内部图如下所示

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

这对于在一次模型通过中测试多层的不同量化非常有用。

未来PR的高层次TODO： * 找出更好的方式来命名输出结构 * 返回一个结果数据结构而不是打印出来 * 在docblocks中添加示例

Return type: 图模块

torch.ao.ns._numeric_suite_fx.loggers_set_enabled(model, enabled)[源代码]¶

在模型的记录器上设置启用设置

torch.ao.ns._numeric_suite_fx.loggers_set_save_activations(model, save_activations)[源代码]¶

设置 model 的日志记录器上的 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。

Return type: 图模块

torch.ao.ns._numeric_suite_fx.extract_results_n_shadows_model(model)[源代码]¶

从模型中提取日志记录结果。

Return type: 字典[字符串, 字典[字符串, 字典[字符串, 列表[字典[字符串, 任意]]]]]

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)[源代码]¶