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='')[source]¶
捕获中间值的基类。
- class torch.ao.ns._numeric_suite_fx.OutputComparisonLogger(*args, **kwargs)[source]¶
与 OutputLogger 相同,但还需要原始激活才能在校准时计算比较
- class torch.ao.ns._numeric_suite_fx.NSTracer(skipped_module_names, skipped_module_classes)[source]¶
就像一个普通的 FX 量化跟踪器一样,但将观察器和 fake_quantize 模块视为叶模块。
- 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)[source]¶
从模型 A 和模型 B 中提取权重,并返回比较结果。
- 参数
model_name_a (str) – 在结果中使用的模型 A 的字符串名称
model_a (Module) – 模型 A
model_name_b (str) – 在结果中使用的模型 B 的字符串名称
model_b (Module) – 模型 B
base_name_to_sets_of_related_ops (Optional[Dict[str, Set[Union[Callable, str]]]]) – 子图基本节点的可选覆盖,可能会更改
unmatchable_types_map (可选[字典[字符串, 集合[联合[可调用, 字符串]]]]) – 可选的不可匹配类型的覆盖,可能随时更改
op_to_type_to_weight_extraction_fn (可选[字典[字符串, 字典[可调用, 可调用]]]) – 可选的提取类型权重的函数覆盖,可能随时更改
- 返回
NSResultsType,包含权重比较结果
- 返回类型
- 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 (字符串) – 模型 A 的字符串名称,用于结果
model_a (Module) – 模型 A
name_b (字符串) – 模型 B 的字符串名称,用于结果
model_b (Module) – 模型 B
logger_cls (可调用) – 要使用的日志记录器类
base_name_to_sets_of_related_ops (Optional[Dict[str, Set[Union[Callable, str]]]]) – 子图基本节点的可选覆盖,可能会更改
unmatchable_types_map (可选[字典[字符串, 集合[联合[可调用, 字符串]]]]) – 可选的不可匹配类型的覆盖,可能随时更改
- 返回
返回一个包含 (model_a_with_loggers, model_b_with_loggers) 的元组。对这两个模型进行就地修改。
- 返回类型
- 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 中的所有日志记录器,并提取记录的信息。
- 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 (字符串) – 模型 A 的字符串名称,用于结果
model_a (Module) – 模型 A
name_b (字符串) – 模型 B 的字符串名称,用于结果
model_b (Module) – 模型 B
logger_cls (可调用) – 要使用的日志记录器类
should_log_inputs (布尔值) – 是否记录输入
base_name_to_sets_of_related_ops (Optional[Dict[str, Set[Union[Callable, str]]]]) – 子图基本节点的可选覆盖,可能会更改
unmatchable_types_map (可选[字典[字符串, 集合[联合[可调用, 字符串]]]]) – 可选的不可匹配类型的覆盖,可能随时更改
- 返回类型
- torch.ao.ns._numeric_suite_fx.extract_shadow_logger_info(model_a_shadows_b, logger_cls, model_name_to_use_for_layer_names)[源代码]¶
遍历影子模型中的所有日志记录器,并提取记录的信息。
- torch.ao.ns._numeric_suite_fx.extend_logger_results_with_comparison(results, model_name_1, model_name_2, comparison_fn, comparison_name)[源代码]¶
使用 comparison_fn 将 model_name_2 中记录的值与 model_name_1 中的对应值进行比较。将结果记录在 model_name_2 的结果中,位于 comparison_name 下。对 results 进行就地修改。
- 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:* 找到一种更好的方法来命名输出结构 * 返回结果数据结构,而不是打印出来 * 将示例添加到文档块中
- 返回类型
- torch.ao.ns._numeric_suite_fx.loggers_set_save_activations(model, save_activations)[source]¶
在 model 的记录器上设置 save_activations 设置。
- torch.ao.ns._numeric_suite_fx.convert_n_shadows_model(model, custom_convert_fn=None, custom_convert_kwargs=None)[source]¶
给定来自 prepare_n_shadows_model 的模型,在每个影子子模块上运行 convert_fx。
- 返回类型