fx quant: add docblocks to _find_matches and _find_quants (#43928)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/43928
Improving readability, no logic change.
Test Plan:
CI
Imported from OSS
Reviewed By: jerryzh168
Differential Revision: D23440249
fbshipit-source-id: a7ebfc7ad15c73e26b9a94758e7254413cc17d29
diff --git a/torch/quantization/fx/quantize.py b/torch/quantization/fx/quantize.py
index 1babf28..fc1a7b6 100644
--- a/torch/quantization/fx/quantize.py
+++ b/torch/quantization/fx/quantize.py
@@ -543,7 +543,27 @@
return quantized
def _find_matches(self, graph, modules, patterns):
- match_map = {} # node name -> (root_node, match_value?)
+ """
+ Matches the nodes in the input graph to quantization patterns, and
+ outputs the information needed to quantize them in future steps.
+
+ Inputs:
+ - graph: an fx.Graph object
+ - modules: a mapping of fully qualified module name to instance,
+ for example, {'foo': ModuleFoo, ...}
+ - patterns: a mapping from a tuple of nodes in reverse order to
+ uninitialized QuantizeHandler subclass.
+
+ Outputs a map of
+ node_name ->
+ (node, matched_values, QuantizeHandler instance, qconfig)
+
+ For example, {
+ 'relu_1': (relu_1, [relu_1], <CopyNode instance>, QConfig(...)),
+ ...
+ }
+ """
+ match_map = {}
all_matched = set()
def record_match(pattern, node, matched):
@@ -570,6 +590,17 @@
return match_map
def _find_quants(self, graph, matches):
+ """
+ Takes the nodes in the input graph and pending matches, and finds and
+ returns the input and output nodes which need to be quantized.
+
+ Inputs:
+ - graph: an fx.Graph object
+ - matches: output of self._find_matches function
+
+ Outputs a map of
+ node_name -> (QuantizeHandler instance (always DefaultQuant), qconfig)
+ """
quants = {}
def visit(node, qconfig):
