Model Speedup Refactor

docs/en_US/Compression/CompressionReference.rst

@@ -137,9 +137,6 @@ Topology Utilities
 ..  autoclass:: nni.compression.pytorch.utils.shape_dependency.GroupDependency
     :members:
 
-..  autoclass:: nni.compression.pytorch.utils.mask_conflict.CatMaskPadding
-    :members:
-
 ..  autoclass:: nni.compression.pytorch.utils.mask_conflict.GroupMaskConflict
     :members:
 

docs/zh_CN/Compression/CompressionReference.rst

@@ -123,9 +123,6 @@ Quantizers
 ..  autoclass:: nni.compression.pytorch.utils.shape_dependency.GroupDependency
     :members:
 
-..  autoclass:: nni.compression.pytorch.utils.mask_conflict.CatMaskPadding
-    :members:
-
 ..  autoclass:: nni.compression.pytorch.utils.mask_conflict.GroupMaskConflict
     :members:
 

nni/common/graph_utils.py

@@ -247,6 +247,7 @@ class TorchModuleGraph(TorchGraph):
     def __init__(self, model=None, dummy_input=None, traced_model=None):
         super().__init__(model, dummy_input, traced_model)
         self.global_count = 0
+        self.reused_module = set()
         self.name_to_node, self.input_to_node, self.output_to_node = self._build_graph()
         self._extract_auxiliary_info()
 
@@ -388,8 +389,14 @@ def _expand_module_node(self, node, node_name, unique_name, op_type, nodes,
                                 visited.add(successor_node)
                         else:
                             outputs.append(output_name)
+                            # if(len(outputs)>1):
+                            #     import pdb; pdb.set_trace()
+                            break
                 else:
                     outputs.append(output_name)
+                    # if(len(outputs)>1):
+                    #     import pdb; pdb.set_trace()
+                # print(_output)
 
         nodepy = NodePyGroup(node_name, unique_name, module_type, op_type,
                              node_group, inputs=list(inputs), outputs=list(outputs))
@@ -724,6 +731,8 @@ def _build_graph(self):
                     unique_name = module_name
                     if use_count > 0:
                         unique_name = module_name + '.%d' % use_count
+                        self.reused_module.add(unique_name)
+                        self.reused_module.add(module_name)
                     node_group = self._expand_module_node(
                         node, module_name, unique_name, module_to_type[module_name],
                         node_cpps, input_to_node, output_to_node, 'module')