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/fuse.py (2998B)

from typing import Dict, Any

from torch.fx import (
    GraphModule,
    Node,
    map_arg
)

from torch.fx.graph import Graph

from ..utils import (
    get_combined_dict
)

from .pattern_utils import (
    get_default_fusion_patterns,
)

from .match_utils import is_match

from .graph_module import (
    FusedGraphModule
)

from .fusion_patterns import *  # noqa: F401,F403

from .quantization_types import Pattern

from typing import Callable, Tuple

class Fuser:
    def fuse(self, model: GraphModule,
             fuse_custom_config_dict: Dict[str, Any] = None) -> GraphModule:
        if fuse_custom_config_dict is None:
            fuse_custom_config_dict = {}

input_root = model
        input_graph = model.graph
        self.modules = dict(input_root.named_modules())

additional_fusion_patterns = \
            fuse_custom_config_dict.get("additional_fusion_pattern", {})
        fusion_patterns = get_combined_dict(
            get_default_fusion_patterns(), additional_fusion_patterns)
        # find fusion
        fusion_pairs = self._find_matches(
            input_root, input_graph, fusion_patterns)
        self.fused_graph = Graph()
        env: Dict[Any, Any] = {}

def load_arg(a):
            return map_arg(a, lambda node: env[node.name])

for node in input_graph.nodes:
            root_node, obj = fusion_pairs.get(node.name, (None, None))
            if root_node is node:
                assert obj is not None
                env[node.name] = obj.fuse(self, load_arg)
            elif root_node is None:
                env[node.name] = self.fused_graph.node_copy(node, load_arg)
            # node matched in patterns and is not root is removed here

preserved_attributes = set(fuse_custom_config_dict.get("preserved_attributes", []))
        model = FusedGraphModule(input_root, self.fused_graph, preserved_attributes)
        return model

def _find_matches(
            self, root: GraphModule, graph: Graph,
            patterns: Dict[Pattern, Callable]
    ) -> Dict[str, Tuple[Node, FuseHandler]]:
        modules = dict(root.named_modules())
        match_map : Dict[str, Tuple[Node, FuseHandler]] = {}  # node name -> (root_node, match_value)

def apply_match(pattern, node, match):
            if isinstance(pattern, tuple):
                s, *args = pattern
                apply_match(s, node, match)
                for subpattern, arg in zip(args, node.args):
                    apply_match(subpattern, arg, match)
            else:
                # the first pattern matches will take precedence
                if node.name not in match_map:
                    match_map[node.name] = match

for node in reversed(graph.nodes):
            if node.name not in match_map:
                for pattern, value in patterns.items():
                    if is_match(modules, node, pattern):
                        apply_match(pattern, node, (node, value(self, node)))

return match_map

from typing import Dict, Any from torch.fx import ( GraphModule, Node, map_arg ) from torch.fx.graph import Graph from ..utils import ( get_combined_dict ) from .pattern_utils import ( get_default_fusion_patterns, ) from .match_utils import is_match from .graph_module import ( FusedGraphModule ) from .fusion_patterns import * # noqa: F401,F403 from .quantization_types import Pattern from typing import Callable, Tuple class Fuser: def fuse(self, model: GraphModule, fuse_custom_config_dict: Dict[str, Any] = None) -> GraphModule: if fuse_custom_config_dict is None: fuse_custom_config_dict = {} input_root = model input_graph = model.graph self.modules = dict(input_root.named_modules()) additional_fusion_patterns = \ fuse_custom_config_dict.get("additional_fusion_pattern", {}) fusion_patterns = get_combined_dict( get_default_fusion_patterns(), additional_fusion_patterns) # find fusion fusion_pairs = self._find_matches( input_root, input_graph, fusion_patterns) self.fused_graph = Graph() env: Dict[Any, Any] = {} def load_arg(a): return map_arg(a, lambda node: env[node.name]) for node in input_graph.nodes: root_node, obj = fusion_pairs.get(node.name, (None, None)) if root_node is node: assert obj is not None env[node.name] = obj.fuse(self, load_arg) elif root_node is None: env[node.name] = self.fused_graph.node_copy(node, load_arg) # node matched in patterns and is not root is removed here preserved_attributes = set(fuse_custom_config_dict.get("preserved_attributes", [])) model = FusedGraphModule(input_root, self.fused_graph, preserved_attributes) return model def _find_matches( self, root: GraphModule, graph: Graph, patterns: Dict[Pattern, Callable] ) -> Dict[str, Tuple[Node, FuseHandler]]: modules = dict(root.named_modules()) match_map : Dict[str, Tuple[Node, FuseHandler]] = {} # node name -> (root_node, match_value) def apply_match(pattern, node, match): if isinstance(pattern, tuple): s, *args = pattern apply_match(s, node, match) for subpattern, arg in zip(args, node.args): apply_match(subpattern, arg, match) else: # the first pattern matches will take precedence if node.name not in match_map: match_map[node.name] = match for node in reversed(graph.nodes): if node.name not in match_map: for pattern, value in patterns.items(): if is_match(modules, node, pattern): apply_match(pattern, node, (node, value(self, node))) return match_map