[Feat] Update ResultCheck

Author: loopyme

src/SubgraphDetection/Util/check_result.py

@@ -2,7 +2,6 @@
 
 from mindinsight.datavisual.data_transform.graph import MSGraph
 
-from SubgraphDetection.DataStructure import SMSGraph
 from SubgraphDetection.DataStructure import Subgraph
 
 
@@ -15,13 +14,7 @@ def __init__(self, graph: MSGraph, subgraph_deque: Deque[Subgraph]):
             subgraph_deque: The result
         """
 
-        self.graph_size = len(
-            tuple(
-                node
-                for node in graph._normal_node_map.values()
-                if node.type not in set(SMSGraph.non_normal_node_type[:1])
-            )
-        )
+        self.graph_size = len(graph._normal_node_map)
         self.num_subgraph = len(subgraph_deque)
         self.subgraph_size = tuple(len(g.nodes[0]) for g in subgraph_deque)
         self.subgraph_count = tuple(len(g.nodes) for g in subgraph_deque)

src/SubgraphDetection/__main__.py

@@ -1,5 +1,5 @@
 import time
-from typing import Deque
+from typing import Deque, Union
 
 from SubgraphDetection import __version__
 from SubgraphDetection.DataStructure import Subgraph
@@ -9,7 +9,7 @@
 from SubgraphDetection.config import CONFIG
 
 
-def detect_subgraph(graph_path, result_path, **kwargs) -> Deque[Subgraph]:
+def detect_subgraph(graph_path, result_path, **kwargs) -> Union[Deque[Subgraph], ResultCheck]:
     """
     Detect the subgraph in a mindspore computational graph
 
@@ -22,7 +22,8 @@ def detect_subgraph(graph_path, result_path, **kwargs) -> Deque[Subgraph]:
         **kwargs: Any other args will pass to config
 
     Returns:
-        Deque of subgraph, all the detected subgraphs
+        Deque of subgraph, all the detected subgraphs.
+        If CHECK_RESULT, it will return a ResultCheck object.
     """
     time_st = time.time()
     CONFIG.set(kwargs)
@@ -31,13 +32,16 @@ def detect_subgraph(graph_path, result_path, **kwargs) -> Deque[Subgraph]:
     result = executor.run()
     dump_result(result, result_path)
 
-    if CONFIG.CHECK_RESULT:
-        print(ResultCheck(graph, result))
     print(
         f"Detecting finished and result have been write to {result_path}, "
         f"total usage of time:{time.time() - time_st} s"
     )
-    return result
+
+    if CONFIG.CHECK_RESULT:
+        result_check = ResultCheck(graph, result)
+        return result_check
+    else:
+        return result
 
 
 def detect_subgraph_in_console():