torch.ao.ns._numeric_suite_fx¶
Warning
This module is an early prototype and is subject to change.
This module contains tooling to compare weights and activations across models. Example usage:
import copy
import torch
import torch.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.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)[source]¶ Base class for capturing intermediate values.
-
class
torch.ao.ns._numeric_suite_fx.NSTracer(skipped_module_names, skipped_module_classes)[source]¶ Just like a regular FX quantization tracer, but treats observers and fake_quantize modules as leaf modules.
-
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]¶ Extract weights from model A and model B, and return a comparison.
- Parameters
model_name_a – string name of model A to use in results
model_a – model A
model_name_b – string name of model B to use in results
model_b – model B
base_name_to_sets_of_related_ops – optional override of subgraph base nodes, subject to change
unmatchable_types_map – optional override of unmatchable types, subject to change
op_to_type_to_weight_extraction_fn – optional override of function which extracts weight from a type, subject to change
- Returns
NSResultsType, containing the weight comparisons
-
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)[source]¶ Instrument model A and model B with loggers.
- Parameters
model_name_a – string name of model A to use in results
model_a – model A
model_name_b – string name of model B to use in results
model_b – model B
logger_cls – class of Logger to use
base_name_to_sets_of_related_ops – optional override of subgraph base nodes, subject to change
unmatchable_types_map – optional override of unmatchable types, subject to change
- Returns
Returns a tuple of (model_a_with_loggers, model_b_with_loggers). Modifies both models inplace.
-
torch.ao.ns._numeric_suite_fx.extract_logger_info(model_a, model_b, logger_cls, model_name_to_use_for_layer_names)[source]¶ Traverse all loggers in model_a and model_b, and extract the logged information.
- Parameters
model_a – model A
model_b – model B
logger_cls – class of Logger to use
model_name_to_use_for_layer_names – string name of model to use for layer names in the output
- Returns
NSResultsType, containing the logged comparisons
-
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)[source]¶ Instrument model A and model B with shadow loggers.
- Parameters
model_name_a – string name of model A to use in results
model_a – model A
model_name_b – string name of model B to use in results
model_b – model B
logger_cls – class of Logger to use
should_log_inputs – whether to log inputs
base_name_to_sets_of_related_ops – optional override of subgraph base nodes, subject to change
unmatchable_types_map – optional override of unmatchable types, subject to change
-
torch.ao.ns._numeric_suite_fx.extract_shadow_logger_info(model_a_shadows_b, logger_cls, model_name_to_use_for_layer_names)[source]¶ Traverse all loggers in a shadow model, and extract the logged information.
- Parameters
model_a_shadows_b – shadow model
logger_cls – class of Logger to use
model_name_to_use_for_layer_names – string name of model to use for layer names in the output
- Returns
NSResultsType, containing the logged comparisons
-
torch.ao.ns._numeric_suite_fx.extend_logger_results_with_comparison(results, model_name_1, model_name_2, comparison_fn, comparison_name)[source]¶ Compares the logged values from model_name_2 against the corresponding values in model_name_1, using comparison_fn. Records the result in model_name_2’s results under comparison_name. Modifies results inplace.
- Parameters
results – the result data structure from extract_logger_info or extract_shadow_logger_info.
model_name_1 – string name of model 1
model_name_2 – string name of model 2
comparison_fn – function to compare two Tensors
model_name_to_use_for_layer_names – string name of model to use for layer names in the output
