Support linear/BN fusion and follow the API guideline (#141585)

jaewoosong · pytorchmergebot · commit 43a2a231d39b · 2024-11-27T06:52:00.000Z
Current `fuse` function supports conv/BN fusions only. This commit is to support linear/BN fusion as well. Changes to follow the API guidelines are also applied. (This will close the PR #141352 which I created for the same topic and got approval but had lint and API guideline problems.) Pull Request resolved: #141585 Approved by: https://github.com/ezyang
diff --git a/torch/fx/experimental/optimization.py b/torch/fx/experimental/optimization.py
@@ -14,7 +14,23 @@
 import torch.utils.mkldnn as th_mkldnn
 from torch.fx.node import Argument, Target
 from torch.fx.passes.shape_prop import ShapeProp
-from torch.nn.utils.fusion import fuse_conv_bn_eval
+from torch.nn.utils.fusion import fuse_conv_bn_eval, fuse_linear_bn_eval
+
+
+__all__ = [
+    "matches_module_pattern",
+    "replace_node_module",
+    "fuse",
+    "remove_dropout",
+    "extract_subgraph",
+    "modules_to_mkldnn",
+    "reset_modules",
+    "MklSubgraph",
+    "gen_mkl_autotuner",
+    "use_mkl_length",
+    "UnionFind",
+    "optimize_for_inference",
+]
 
 
 def _parent_name(target: str) -> Tuple[str, str]:
@@ -58,13 +74,14 @@ def replace_node_module(
 
 def fuse(model: torch.nn.Module, inplace=False, no_trace=False) -> torch.nn.Module:
     """
-    Fuses convolution/BN layers for inference purposes. Will deepcopy your
-    model by default, but can modify the model inplace as well.
+    Fuses convolution/BN and linear/BN layers for inference purposes.
+    Will deepcopy your model by default, but can modify the model inplace as well.
     """
     patterns = [
         (nn.Conv1d, nn.BatchNorm1d),
         (nn.Conv2d, nn.BatchNorm2d),
         (nn.Conv3d, nn.BatchNorm3d),
+        (nn.Linear, nn.BatchNorm1d),
     ]
     if not inplace:
         model = copy.deepcopy(model)
@@ -78,14 +95,18 @@ def fuse(model: torch.nn.Module, inplace=False, no_trace=False) -> torch.nn.Modu
     for pattern in patterns:
         for node in new_graph.nodes:
             if matches_module_pattern(pattern, node, modules):
-                if len(node.args[0].users) > 1:  # Output of conv is used by other nodes
+                if len(node.args[0].users) > 1:
+                    # Output of conv/linear is used by other nodes
                     continue
-                conv = modules[node.args[0].target]
+                first_layer = modules[node.args[0].target]
                 bn = modules[node.target]
                 if not bn.track_running_stats:
                     continue
-                fused_conv = fuse_conv_bn_eval(conv, bn)
-                replace_node_module(node.args[0], modules, fused_conv)
+                if pattern[0] in [nn.Conv1d, nn.Conv2d, nn.Conv3d]:
+                    fused_layer = fuse_conv_bn_eval(first_layer, bn)
+                else:  # nn.Linear
+                    fused_layer = fuse_linear_bn_eval(first_layer, bn)
+                replace_node_module(node.args[0], modules, fused_layer)
                 node.replace_all_uses_with(node.args[0])
                 new_graph.erase_node(node)
     return fx.GraphModule(fx_model, new_graph)
