torch.ao.ns._numeric_suite_fx

創建於: 2025年6月13日 | 最後更新於: 2025年6月13日

警告

This module is an early prototype and is subject to change.

此模組包含用於比較模型之間權重和啟用的工具。用法示例

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 的字串名稱

• 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，包含權重比較

返回型別

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)

為 model A 和 model B 新增 logger 儀器。

引數

• name_a (str) – model A 在結果中使用的字串名稱

• model_a (Module) – 模型 A

• name_b (str) – model 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 中的所有 logger，並提取記錄的資訊。

引數

• model_a (Module) – 模型 A

• model_b (Module) – 模型 B

• 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.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)

為 model A 和 model B 新增 shadow logger 儀器。

引數

• name_a (str) – model A 在結果中使用的字串名稱

• model_a (Module) – 模型 A

• name_b (str) – model B 在結果中使用的字串名稱

• model_b (Module) – 模型 B

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

• 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)

遍歷 shadow model 中的所有 logger，並提取記錄的資訊。

引數

• model_a_shadows_b (Module) – shadow model

• 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)

使用 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) – model 1 的字串名稱

• model_name_2 (str) – model 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 個操作的圖的模型，例如

args_kwargs_m -> op_m -> output_m

以及每個 op 的 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 ---

這對於在單次模型透過中測試多個層的不同量化非常有用。

高階 TODO 列表，用於後續 PR： * 找到一種更好的方法來命名輸出結構 * 返回結果資料結構而不是列印它 * 向 docblocks 新增示例

返回型別

GraphModule

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

設定 model 的 logger 的 enabled 設定

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

設定 model 的 logger 的 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 的模型，對每個 shadow 子模組執行 convert_fx。

返回型別

GraphModule

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

從 model 中提取 logger 結果。

返回型別

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

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

列印提取的 results 的摘要。

---

torch.ao.ns.fx.utils

警告

This module is an early prototype and is subject to change.

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)