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.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='')

Base class for capturing intermediate values.

forward(x)

class torch.ao.ns._numeric_suite_fx.OutputComparisonLogger(*args, **kwargs)

Same as OutputLogger, but also requires the original activation in order to calculate the comparison at calibration time

forward(x, x_ref)

class torch.ao.ns._numeric_suite_fx.NSTracer(skipped_module_names, skipped_module_classes)

Just like a regular FX quantization tracer, but treats observers and fake_quantize modules as leaf modules.

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)

Extract weights from model A and model B, and return a comparison.

Parameters:

• model_name_a (str) – string name of model A to use in results

• model_a (Module) – model A

• model_name_b (str) – string name of model B to use in results

• model_b (Module) – model B

• base_name_to_sets_of_related_ops (Optional[Dict[str, Set[Union[Callable, str]]]]) – optional override of subgraph base nodes, subject to change

• unmatchable_types_map (Optional[Dict[str, Set[Union[Callable, str]]]]) – optional override of unmatchable types, subject to change

• op_to_type_to_weight_extraction_fn (Optional[Dict[str, Dict[Callable, Callable]]]) – optional override of function which extracts weight from a type, subject to change

Returns:

NSResultsType, containing the weight comparisons

Return type:

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)

Instrument model A and model B with loggers.

Parameters:

• name_a (str) – string name of model A to use in results

• model_a (Module) – model A

• name_b (str) – string name of model B to use in results

• model_b (Module) – model B

• logger_cls (Callable) – class of Logger to use

• base_name_to_sets_of_related_ops (Optional[Dict[str, Set[Union[Callable, str]]]]) – optional override of subgraph base nodes, subject to change

• unmatchable_types_map (Optional[Dict[str, Set[Union[Callable, str]]]]) – 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.

Return type:

Tuple[Module, Module]

torch.ao.ns._numeric_suite_fx.extract_logger_info(model_a, model_b, logger_cls, model_name_to_use_for_layer_names)

Traverse all loggers in model_a and model_b, and extract the logged information.

Parameters:

• model_a (Module) – model A

• model_b (Module) – model B

• logger_cls (Callable) – class of Logger to use

• model_name_to_use_for_layer_names (str) – string name of model to use for layer names in the output

Returns:

NSResultsType, containing the logged comparisons

Return type:

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)

Instrument model A and model B with shadow loggers.

Parameters:

• name_a (str) – string name of model A to use in results

• model_a (Module) – model A

• name_b (str) – string name of model B to use in results

• model_b (Module) – model B

• logger_cls (Callable) – class of Logger to use

• should_log_inputs (bool) – whether to log inputs

• base_name_to_sets_of_related_ops (Optional[Dict[str, Set[Union[Callable, str]]]]) – optional override of subgraph base nodes, subject to change

• unmatchable_types_map (Optional[Dict[str, Set[Union[Callable, str]]]]) – optional override of unmatchable types, subject to change

Return type:

Module

torch.ao.ns._numeric_suite_fx.extract_shadow_logger_info(model_a_shadows_b, logger_cls, model_name_to_use_for_layer_names)

Traverse all loggers in a shadow model, and extract the logged information.

Parameters:

• model_a_shadows_b (Module) – shadow model

• logger_cls (Callable) – class of Logger to use

• model_name_to_use_for_layer_names (str) – string name of model to use for layer names in the output

Returns:

NSResultsType, containing the logged comparisons

Return type:

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)

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 (Dict[str, Dict[str, Dict[str, List[Dict[str, Any]]]]]) – the result data structure from extract_logger_info or extract_shadow_logger_info.

• model_name_1 (str) – string name of model 1

• model_name_2 (str) – string name of model 2

• comparison_fn (Callable[[Tensor, Tensor], Tensor]) – function to compare two Tensors

• comparison_name (str) – string name of model to use for layer names in the output

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)

Given a model with a graph with M ops such as

args_kwargs_m -> op_m -> output_m

And a set of N qconfigs for each op, creates a new model, with each of the subgraph of op_m transformed into

     |---------> op_m_n -> log_m_n
     |                     /
args_kwargs_m ---------> op_m -> log_m_0

Where op_m_n is op_m wrapped in a submodule and transformed with qconfig_n, and its inner graph looks like

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

This is useful for testing different quantization of multiple layers in a single pass through the model.

High level TODOs for future PRs: * figure out a better way to name the output structure * return a results data structure instead of printing it out * add examples to docblocks

Return type:

GraphModule

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

Sets the enabled setting on a model’s loggers

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

Sets the save_activations setting on a model’s loggers

torch.ao.ns._numeric_suite_fx.convert_n_shadows_model(model, custom_convert_fn=None, custom_convert_kwargs=None)

Given a model from prepare_n_shadows_model, runs convert_fx on each shadow submodule.

Return type:

GraphModule

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

Extracts logger results from model.

Return type:

Dict[str, Dict[str, Dict[str, List[Dict[str, Any]]]]]

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

Prints a summary of extracted results.

torch.ao.ns.fx.utils

Warning

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)