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.
- classtorch.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)[源代码]
- classtorch.ao.ns._numeric_suite_fx.OutputComparisonLogger(*args, **kwargs)[源代码]
-
与OutputLogger类似,但需要原始激活值来在校准时进行比较
- forward(x, x_ref)[源代码]
- classtorch.ao.ns._numeric_suite_fx.NSTracer(skipped_module_names, skipped_module_classes)[源代码]
-
就像一个常规的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)[源代码]
-
从模型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 (Optional[Dict[str, Set[Union[Callable, str]]]]) – 可选的不匹配类型的覆盖,可能会发生变化
-
op_to_type_to_weight_extraction_fn (Optional[Dict[str, Dict[Callable, Callable]]]) – 可选的函数覆盖,用于从类型中提取权重,可能会发生变化
-
- 返回值
-
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 (str) – 模型 A 在结果中的字符串名称
-
model_a (Module) — 模型A
-
name_b (str) – 模型 B 在结果中的字符串名称
-
model_b (Module) – 模型B
-
logger_cls (Callable) – 指定要使用的日志记录器类
-
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) 的元组,并就地修改这两个模型。
- 返回类型
- 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 (str) – 模型 A 在结果中的字符串名称
-
model_a (Module) — 模型A
-
name_b (str) – 模型 B 在结果中的字符串名称
-
model_b (Module) – 模型B
-
logger_cls (Callable) – 指定要使用的日志记录器类
-
should_log_inputs (bool) – 是否记录输入数据
-
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)[源代码]
-
遍历影子模型中的所有日志记录器,并提取日志信息。
- 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。
- 参数
-
-
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]) – 用于比较两个 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个操作(ops)的模型图,例如:
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的高层次待办事项:* 确定更佳的输出结构命名方式 * 返回结果数据结构,而非直接打印 * 在文档块中增加示例
- 返回类型
- torch.ao.ns._numeric_suite_fx.loggers_set_enabled(model, enabled)[源代码]
-
设置模型的日志记录器的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。
- 返回类型
- torch.ao.ns._numeric_suite_fx.extract_results_n_shadows_model(model)[源代码]
-
从模型中提取日志记录结果。
- torch.ao.ns._numeric_suite_fx.print_comparisons_n_shadows_model(results)[源代码]
-
打印结果的摘要。