FM / fx

/usr/local/lib64/python3.6/site-packages/torch/quantization/fx/

/usr/local/lib64/python3.6/site-packages/torch/quantization/fx

Name	Size	Mode	Actions
backend_config_dict/	-	0755	rm
__pycache__/	-	0755	rm
convert.py	24481	0644	edit dl rm
fuse.py	2998	0644	edit dl rm
fusion_patterns.py	6862	0644	edit dl rm
graph_module.py	5117	0644	edit dl rm
lower_to_fbgemm.py	337	0644	edit dl rm
lower_to_qnnpack.py	339	0644	edit dl rm
match_utils.py	6278	0644	edit dl rm
pattern_utils.py	2171	0644	edit dl rm
prepare.py	48970	0644	edit dl rm
qconfig_utils.py	11818	0644	edit dl rm
quantization_patterns.py	97807	0644	edit dl rm
quantization_types.py	319	0644	edit dl rm
quantized_fusion_patterns_and_replacements.py	876	0644	edit dl rm
utils.py	20698	0644	edit dl rm
_equalize.py	36240	0644	edit dl rm
_lower_to_native_backend.py	741	0644	edit dl rm
__init__.py	212	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/quantization/fx/pattern_utils.py (2171B)

# TODO(future PR): fix the typing on QuantizeHandler (currently a circular dependency)
QuantizeHandler = Any

MatchResult = Tuple[Node, List[Node], Optional[Pattern], QuantizeHandler,
                    QConfigAny]

# pattern for conv bn fusion
DEFAULT_FUSION_PATTERNS = OrderedDict()
def register_fusion_pattern(pattern):
    def insert(fn):
        DEFAULT_FUSION_PATTERNS[pattern] = fn
        return fn
    return insert

def get_default_fusion_patterns() -> Dict[Pattern, QuantizeHandler]:
    return DEFAULT_FUSION_PATTERNS

DEFAULT_QUANTIZATION_PATTERNS = OrderedDict()
# a map from pattern to activation_post_process(observer/fake_quant) consstructor for output activation
# e.g. pattern: torch.sigmoid,
#      output_activation_post_process: default_affine_fixed_qparam_fake_quant
DEFAULT_OUTPUT_ACTIVATION_POST_PROCESS_MAP = dict()

# Register pattern for both static quantization and qat
def register_quant_pattern(pattern, output_activation_post_process=None):
    def insert(fn):
        DEFAULT_QUANTIZATION_PATTERNS[pattern] = fn
        if output_activation_post_process is not None:
            DEFAULT_OUTPUT_ACTIVATION_POST_PROCESS_MAP[pattern] = output_activation_post_process
        return fn
    return insert

# Get patterns for both static quantization and qat
def get_default_quant_patterns() -> Dict[Pattern, QuantizeHandler]:
    return DEFAULT_QUANTIZATION_PATTERNS

# a map from pattern to output activation post process constructor
# e.g. torch.sigmoid -> default_affine_fixed_qparam_fake_quant
def get_default_output_activation_post_process_map() -> Dict[Pattern, torch.quantization.observer.ObserverBase]:
    return DEFAULT_OUTPUT_ACTIVATION_POST_PROCESS_MAP

# Example use of register pattern function:
# @register_fusion_pattern(torch.nn.ReLU, (torch.nn.BatchNorm2d, torch.nn.Conv2d)))
# class ConvBNReLUFusion():
#     def __init__(...):
#         ...
#

import torch from collections import OrderedDict from typing import Dict, Any, Tuple, List, Optional from torch.fx.graph import ( Node, ) from .quantization_types import Pattern from ..qconfig import QConfigAny # from .quantization_patterns import BinaryOpQuantizeHandler # TODO(future PR): fix the typing on QuantizeHandler (currently a circular dependency) QuantizeHandler = Any MatchResult = Tuple[Node, List[Node], Optional[Pattern], QuantizeHandler, QConfigAny] # pattern for conv bn fusion DEFAULT_FUSION_PATTERNS = OrderedDict() def register_fusion_pattern(pattern): def insert(fn): DEFAULT_FUSION_PATTERNS[pattern] = fn return fn return insert def get_default_fusion_patterns() -> Dict[Pattern, QuantizeHandler]: return DEFAULT_FUSION_PATTERNS DEFAULT_QUANTIZATION_PATTERNS = OrderedDict() # a map from pattern to activation_post_process(observer/fake_quant) consstructor for output activation # e.g. pattern: torch.sigmoid, # output_activation_post_process: default_affine_fixed_qparam_fake_quant DEFAULT_OUTPUT_ACTIVATION_POST_PROCESS_MAP = dict() # Register pattern for both static quantization and qat def register_quant_pattern(pattern, output_activation_post_process=None): def insert(fn): DEFAULT_QUANTIZATION_PATTERNS[pattern] = fn if output_activation_post_process is not None: DEFAULT_OUTPUT_ACTIVATION_POST_PROCESS_MAP[pattern] = output_activation_post_process return fn return insert # Get patterns for both static quantization and qat def get_default_quant_patterns() -> Dict[Pattern, QuantizeHandler]: return DEFAULT_QUANTIZATION_PATTERNS # a map from pattern to output activation post process constructor # e.g. torch.sigmoid -> default_affine_fixed_qparam_fake_quant def get_default_output_activation_post_process_map() -> Dict[Pattern, torch.quantization.observer.ObserverBase]: return DEFAULT_OUTPUT_ACTIVATION_POST_PROCESS_MAP # Example use of register pattern function: # @register_fusion_pattern(torch.nn.ReLU, (torch.nn.BatchNorm2d, torch.nn.Conv2d))) # class ConvBNReLUFusion(): # def __init__(...): # ... #