Table of Contents
Shortcuts

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

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.