mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

mindspore/rewrite/symbol_tree.py

@@ -14,30 +14,31 @@
 # ============================================================================
 """SymbolTree class define of Rewrite according to forward function of a network."""
 import stat
-from typing import Optional, Union, Tuple, Any
+from typing import Optional, Union, Tuple, Any, Dict, List
 import os
 import sys
 import ast
-import importlib
-import types
+import importlib.util
 import time
-import astunparse
 
 from mindspore.nn import Cell
 from mindspore import log as logger
-from mindspore.rewrite.ast_creator_register import ast_creator_registry
-from .node import Node, TreeNode
+from .node.node import Node, TreeNode
 from .api.node_type import NodeType
-from .ast_helpers import AstModifier, AstReplacer, StrChecker, AstFinder, CheckPropertyIsUsed
+from .ast_helpers import AstModifier, AstReplacer, StrChecker, AstFinder, AstClassFinder, AstFunctionFinder
 from .api.scoped_value import ScopedValue, ValueType
 from .symbol_tree_dumper import SymbolTreeDumper
-from .topological_manager import TopoManager
+from .node.node_topological_manager import TopoManager
 from .namer import TargetNamer, NodeNamer, ClassNamer
 from .common.observer import Observer
 from .common.observable import Observable
 from .common.event import Event
-from .node_visitor import NodeVisitor
+from .node.node_manager import NodeManager
 
+if sys.version_info >= (3, 9):
+    import ast as astunparse # pylint: disable=reimported, ungrouped-imports
+else:
+    import astunparse
 
 class Position:
     """
@@ -81,6 +82,7 @@ class FieldFinder(AstFinder):
     Args:
         scope (ast.AST): An instance of ast node as search scope.
     """
+
     def __init__(self, scope: ast.AST):
         super().__init__(scope)
         self._result = False
@@ -134,7 +136,7 @@ class IfFixer(ast.NodeTransformer):
         self.generic_visit(node)
 
 
-class SymbolTree(Observer, Observable):
+class SymbolTree(Observer, Observable, NodeManager):
     """
     A symbol-tree usually corresponding to forward method of a network.
 
@@ -147,227 +149,138 @@ class SymbolTree(Observer, Observable):
     """
 
     def __init__(self, origin_network: Cell, module_ast: ast.Module):
-        super().__init__()
+        Observer.__init__(self)
         Observable.__init__(self)
-        origin_network_key = "handler"
+        self._node_namer = NodeNamer()
+        self._node_namer.add_name('obj')
+        NodeManager.__init__(self, self._node_namer)
+        NodeManager.reg_observer(self, observer=self)
         # init unique-namers
         self._target_namer = TargetNamer()
-        self._node_name_namer = NodeNamer()
-        # name or node would use as name of field, so name of origin network handler field should be added into \
-        # _node_name_namer.
-        self._node_name_namer.add_name(origin_network_key)
-        self._topo_mgr = TopoManager()
-        self._topo_mgr.reg_observer(self)
-
-        self._nodes: {str, Node} = {}
-        # parameters of forward method
-        self._inputs: [Node] = []
+        # input arguments of function
         self._ori_cls_name = type(origin_network).__name__
         self._opt_cls_name = ClassNamer.instance().get_name(self._ori_cls_name)
+        NodeManager.set_manager_name(self, self._opt_cls_name)
         self._origin_network = origin_network
         self._module_ast: ast.Module = module_ast
+        self._import_asts: Optional[ast.Ast] = []
         self._class_ast: Optional[ast.ClassDef] = None
         self._root_ast: Optional[ast.FunctionDef] = None
         self._init_func_ast: Optional[ast.FunctionDef] = None
         self._deleted_field = {}
         self._deleted_node = []
-        self._external_func_ast = []
+        self._external_ast = []
         self._father_class_ast = []
-
-        # head node is always point to the first node(in source code order) of SymbolTree
-        self._head = None
-        # tail node is always point to the last node(in source code order) of SymbolTree
-        self._tail = None
-        self._return: Optional[Node] = None
-
         self._modified = False
-        self._node_visitor = None
-
         self._tmp_file_limits = 20
         self._tmp_files = []
         self._saved_file_name = "./network_define.py"
+        # used to insert "sys.path.append(xxx)"
+        self._net_file_paths = []
+        self._tmp_import_strs = []
+        self._tmp_unmodified_strees: {type, str} = {}
+        self._tmp_replacers = []
+        # Record imported modules and names of each files
+        # The meanings of `module` and `name` are like code: from `module` import `nameA`, `nameB`
+        # Format: {file_path: {module: [name, ...], ...}, ...}
+        self._imported_modules: Dict[str, Dict[str, List[str]]] = {}
 
     def __del__(self):
         for tmp_file in self._tmp_files:
             tmp_file.close()
 
     @staticmethod
-    def _find_consumers_and_providers(nodes: [Node]):
-        """
-        Find consumers and providers for all nodes according to their targets and arguments.
-        """
-        consumers: {ScopedValue: [Node]} = {}
-        providers: {ScopedValue: Node} = {}
-        for node in nodes:
-            for arg in node.get_args():
-                if consumers.get(arg):
-                    consumers[arg].append(node)
-                else:
-                    consumers[arg] = [node]
-            for _, arg in node.get_kwargs():
-                if consumers.get(arg):
-                    consumers[arg].append(node)
-                else:
-                    consumers[arg] = [node]
-            for target in node.get_targets():
-                if providers.get(target) is not None:
-                    raise RuntimeError(f"Target({target}) of node duplicated")
-                providers[target] = node
-        return consumers, providers
-
-    @staticmethod
-    def _link_nodes_and_find_root(nodes: [Node]) -> Node:
-        """
-        Find inputs for all nodes created by Replacement according to their targets and arguments.
-
-        Find root node of all nodes created by Replacement. One and Only one root should be found.
-
-        Args:
-            nodes (list[Node]): A list of instance of Node created by Replacement.
-
-        Returns:
-            An instance of Node represents root of input nodes.
-        """
-        consumers, providers = SymbolTree._find_consumers_and_providers(nodes)
-        # find root node
-        root = None
-        for node in nodes:
-            used = 0
-            for target in node.get_targets():
-                cur_consumers = consumers.get(target)
-                if not cur_consumers:
-                    continue
-                for cur_consumer in cur_consumers:
-                    if id(cur_consumer) != id(node):
-                        used += 1
-                        break
-            if used == 0:
-                if root is not None:
-                    raise RuntimeError("Replacement should only has one root, found multi-root")
-                root = node
-        if root is None:
-            raise RuntimeError("Replacement should only has one root, found no root")
-        # link node's input
-        for node in nodes:
-            inputs = []
-            for _, arg in node.get_normalized_args().items():
-                node_input: Node = providers.get(arg)
-                if id(node_input) != id(node):
-                    inputs.append(node_input)
-            node.set_inputs(inputs)
-        return root
-
-    @staticmethod
-    def _find_all_class_in_symboltree(stree: 'SymbolTree', seen_class: {type, str}, allow_class_name: [], replacers):
-        """Find all non-duplicated class name of SymbolTree recursively."""
-        replacer = AstReplacer(stree._class_ast)
-        replacers.append(replacer)
-        for node in stree.nodes():
-            if not isinstance(node, TreeNode):
+    def _remove_unused_import(module_ast):
+        """remove unused import in self._module_ast"""
+        str_checker = StrChecker(module_ast)
+        for i in range(len(module_ast.body) - 1, -1, -1):
+            body = module_ast.body[i]
+            if not isinstance(body, (ast.Import, ast.ImportFrom)):
                 continue
-            if node.symbol_tree._class_ast is None:
+            if isinstance(body, ast.Import):
                 continue
-            sub_stree: SymbolTree = node.symbol_tree
-            SymbolTree._find_all_class_in_symboltree(sub_stree, seen_class, allow_class_name, replacers)
-            # all modified ast.ClassDef should export to code
-            if sub_stree._modified:
-                allow_class_name.append(sub_stree._class_ast.name)
+            if isinstance(body, ast.ImportFrom) and body.module == "cell":
+                module_ast.body.remove(body)
                 continue
-            # all un-modified ast.ClassDef only keep one instance
-            seen_cls_name = seen_class.get(type(sub_stree.get_origin_network()))
-            if seen_cls_name is not None:
-                replacer.replace_all(sub_stree._class_ast.name, seen_cls_name)
-            else:
-                seen_class[type(sub_stree.get_origin_network())] = sub_stree._class_ast.name
-                allow_class_name.append(sub_stree._class_ast.name)
+            for alias in body.names:
+                name = alias.asname if alias.asname else alias.name
+                if not str_checker.check(name):
+                    if len(body.names) == 1:
+                        module_ast.body.remove(body)
+                        i += 1
+                    else:
+                        body.names.remove(alias)
+
+    @staticmethod
+    def _remove_duplicated_import(module_ast):
+        """Remove duplicated import of 'net'."""
+        imports = set()
+        futures = set()
+        classes = set()
+
+        class TransImportNode(ast.NodeTransformer):
+            """Find all import nodes from input ast node."""
+
+            def visit_ClassDef(self, node: ast.ClassDef) -> Any:
+                class_str = astunparse.unparse(node)
+                if class_str not in classes:
+                    classes.add(node.name)
+                    return node
+                return
+
+            def visit_Try(self, node: ast.Try) -> Any:
+                if isinstance(node.body[0], (ast.Import, ast.ImportFrom)):
+                    import_str = astunparse.unparse(node)
+                    if import_str not in imports:
+                        imports.add(import_str)
+                        return node
+                return
+
+            def visit_Import(self, node: ast.Import) -> Any:
+                import_str = astunparse.unparse(node)
+                if import_str not in imports:
+                    imports.add(import_str)
+                    return node
+                return
+
+            def visit_ImportFrom(self, node: ast.ImportFrom) -> Any:
+                """
+                Once the father class 'A' is defined in the current module, all the next imported class 'A' should
+                be removed. e.g.
+                    def class A():
+                        ...
+                    from xxx import A, B
+                    =>
+                    def class A():
+                        ...
+                    from xxx import B
+                """
+                import_str = astunparse.unparse(node)
+
+                if import_str not in imports:
+                    imports.add(import_str)
+                    # remove "__future__" module
+                    if node.module == '__future__':
+                        futures.add(node.module)
+                        return
+                    # remove modules which have been defined in the code file
+                    # it occurs when class A is a father class and other sub-classes import A
+                    for alias in node.names[:]:
+                        if alias.name in classes:
+                            node.names.remove(alias)
+                    # if the alias(es) in node.names are all removed, this import statement should be removed
+                    if not node.names:
+                        return
+                    return node
+                return
+
+        get_node_handler = TransImportNode()
+        get_node_handler.generic_visit(module_ast)
 
     def finish_build(self):
         """Add Event.TopologicalChangeEvent event when build is finished."""
         self.add_event(Event.TopologicalChangeEvent)
 
-    def _create_call_function(self, func, targets, args, kwargs):
-        """
-        Create a Node object and generate the execution code to insert into the source code.
-        The source code calls the 'func' function with 'args' and' kwargs' as parameters.
-
-        Args:
-            func (FunctionType) - The function to be called.
-            targets (list [str]) - indicates the output name. As the output of the node in the source code.
-            args (ParamType) - parameter name of the node. Used as a parameter to a code statement in source
-                code. The default value is None, which means there is no parameter input in the cell.
-            kwargs ({str: ParamType}) - The key type must be str, and the value type must be ParamType. The
-                input parameter name used to describe the formal parameter with a keyword. Enter the name in the source
-                code as the 'kwargs' in the statement expression. The default value is None, which means there is no
-                'kwargs' input.
-
-        Returns:
-            An instance of `Node`.
-        """
-        if not isinstance(func, types.FunctionType):
-            raise TypeError("The 'func' parameter must be a Function, but got ", type(func))
-
-        _package = func.__globals__['__package__']
-        func_name = ".".join([_package, func.__name__]) if _package else func.__name__
-
-        ast_assign = self.create_assign_node(targets, func_name, args, kwargs)
-        scope_targets = [ScopedValue.create_naming_value(targets[0])]
-        scope_func = ScopedValue.create_naming_value(func_name, "")
-        call_args = list()
-        for arg in args:
-            if isinstance(arg, Node):
-                call_args.append(ScopedValue.create_variable_value(arg.get_targets()[0].value))
-            else:
-                call_args.append(ScopedValue.create_variable_value(arg))
-        call_kwargs = {}
-        for k, v in kwargs.items():
-            call_kwargs[k] = ScopedValue.create_variable_value(v)
-        node = self.inner_create_call_function(func_name, ast_assign, scope_func, func, scope_targets, call_args,
-                                               call_kwargs)
-        return node
-
-    def create_assign_node(self, targets, func_name, args, kwargs):
-        """
-        Create a ast.Assign type node.
-
-        Args:
-            targets (list): _description_
-            func_name (_type_): _description_
-            args (_type_): _description_
-            kwargs (_type_): _description_
-
-        Returns:
-            _type_: _description_
-        """
-        # create targets
-        ast_targets = [ast_creator_registry.get("Name")(targets)]
-        # create call
-        ast_func = ast_creator_registry.get("Attribute")(func_name)
-        ast_args = ast_creator_registry.get("Args")(args)
-        ast_kwargs = ast_creator_registry.get("KwArgs")(kwargs) if kwargs else []
-        ast_value = ast_creator_registry.get("Call")(func=ast_func, args=ast_args, keywords=ast_kwargs)
-        # create assign
-        ast_node = ast_creator_registry.get("Assign")(targets=ast_targets, value=ast_value)
-        return ast_node
-
-    def inner_create_call_function(self, node_name, ast_node, func_name, func, targets, args, kwargs):
-        '''
-        Instantiate an instance of node whose type is `CallFunction`.
-
-        Args:
-            node_name (str): Name of node.
-            func_name (str): Name of function.
-            ast_node ([ast.AST, optional]): An instance of ast.AST represents corresponding node in ast.
-            targets (list[ScopedValue]): A list of instance of `ScopedValue`. See detail in docstring of Node class.
-            func ([ScopedValue, optional]): An instance of `ScopedValue`. See detail in docstring of Node class.
-            args (list[ScopedValue]): A list of instance of `ScopedValue`. See detail in docstring of Node class.
-            kwargs (dict{str: ScopedValue}): A list of instance of `ScopedValue`. See detail in docstring of `Node`
-                class.
-        '''
-        logger.info(f"func name: {func_name}; func: {func}; targets: {targets}; args: {args}; kwargs: {kwargs}")
-        node = Node(NodeType.CallFunction, ast_node, targets, func_name, args, kwargs, node_name, func)
-        node.set_belong_symbol_tree(self)
-        return node
-
     def get_ori_cls_name(self) -> str:
         """
         Get class name of original network.
@@ -423,6 +336,7 @@ class SymbolTree(Observer, Observable):
                                         corresponding network class.
         """
         self._root_ast = ast_node
+        NodeManager.set_ast_functiondef(self, ast_node)
 
     def get_class_ast(self):
         """
@@ -461,18 +375,6 @@ class SymbolTree(Observer, Observable):
         """
         self._init_func_ast = ast_node
 
-    def get_inputs(self):
-        return self._inputs
-
-    def get_head_node(self):
-        """
-        Getter of `_head` which represents the beginning node while iterating SymbolTree nodes.
-
-        Returns:
-            An instance of node.
-        """
-        return self._head
-
     def get_origin_network(self):
         """
         Getter of `_origin_network`.
@@ -486,36 +388,53 @@ class SymbolTree(Observer, Observable):
         """Get dict of nodes"""
         return self._nodes
 
-    def nodes(self):
+    def get_node_namer(self):
+        """Get _node_namer"""
+        return self._node_namer
+
+    def is_modified(self):
         """
-        Get generator of nodes of current `SymbolTree`.
+        Check whether symbol tree is modified.
 
-        Returns:
-            A generator for iterating Nodes of `SymbolTree`.
+        Symbol tree is considered as modified if operations like insert/replace/erase/set_arg is called after
+        the symbol tree is created.
         """
-        if self._node_visitor is None:
-            self._node_visitor = NodeVisitor(self)
-        it = iter(self._node_visitor)
+        return self._modified
 
-        while True:
-            try:
-                n = next(it)
-            except StopIteration:
-                return None
-            yield n
+    def set_modified_true(self):
+        """
+        Set self._modified true.
 
-    def get_node(self, node_name: str) -> Optional[Node]:
+        Self._modified is set true when 'if' exists in the original network.
+        In this situation, different original network instance tends to be different.
+        Hence, the class name should be updated.
         """
-        Get node of current symbol_tree by `node_name`.
+        self._modified = True
 
-        Args:
-            node_name (str): A str represents name of node as key of query.
+    def get_import_asts(self):
+        """Get _import_asts"""
+        return self._import_asts
 
-        Returns:
-            An instance of Node if found else None.
-        """
+    def get_external_ast(self):
+        """Get _external_ast"""
+        return self._external_ast
+
+    def get_father_class_ast(self):
+        """Get _father_class_ast"""
+        return self._father_class_ast
 
-        return self._nodes.get(node_name)
+    def get_imported_modules(self, file_path: str):
+        """Get all modules and module_paths in file of `file_path` ."""
+        return self._imported_modules.get(file_path, {})
+
+    def save_imported_modules(self, file_path: str, module: str, names: List[str]):
+        """Save module and names into _imported_modules."""
+        imported_modules = self.get_imported_modules(file_path)
+        if imported_modules.get(module):
+            imported_modules[module].extend(names)
+        else:
+            imported_modules[module] = names
+        self._imported_modules[file_path] = imported_modules
 
     def get_node_inputs(self, node_or_name: Union[Node, str]) -> [Node]:
         """
@@ -553,7 +472,7 @@ class SymbolTree(Observer, Observable):
             return []
         if real_node.get_node_type() == NodeType.Output:
             return []
-        return self._topo_mgr.get_node_users(node_or_name)
+        return TopoManager.get_node_users(real_node)
 
     def before(self, node_or_name: Union[Node, str]) -> Position:
         """
@@ -606,9 +525,11 @@ class SymbolTree(Observer, Observable):
             raise RuntimeError("Node is not belong to current SymbolTree: ", node_or_name)
         return Position.create(node.get_belong_symbol_tree(), node, False)
 
-    def insert_node(self, position: Optional[Position], node: Node, insert_to_ast: bool = True) -> Node:
+    def insert_node(self, new_node: Node, base_node: Node, before_node: bool, node_manager: NodeManager = None,
+                    insert_to_ast: bool = True):
         """
-        Insert a node into SymbolTree.
+        Insert a node before or after base_node.
+
         Note:
             Name of node will be unique while inserting node into SymbolTree.
 
@@ -627,52 +548,73 @@ class SymbolTree(Observer, Observable):
             Topological relation is updated and inputs of corresponding node is updated.
 
         Args:
-            position (Position): A Position indicates an insert position point.
-            node (Node): An instance of node to be inserted in.
-            insert_to_ast (bool): A bool indicates whether to update corresponding ast node at same time, default is
-                True.
+            new_node (Node): Node to be inserted.
+            base_node (Node): New node will be inserted before or after base_node.
+            before_node (bool): Indicate whether new node is inserted before base_node.
+            node_manager (NodeManager): NodeManager those asts belong to. Default: None, means those asts belong to
+                NodeManager of symboltree's construct function.
+            insert_to_ast (bool): Indicate whether ast nodes need to be updated.
 
         Returns:
             An instance of node which has been inserted into SymbolTree.
 
         Raises:
-            RuntimeError: If 'position' is not in current SymbolTree.
+            ValueError: Node in the SymbolTree is inserted into SymbolTree again.
             RuntimeError: If corresponding ast node is not an ast.Assign when 'insert_to_ast' is True.
         """
-        if position is not None and hasattr(position.node, "container"):
-            cellcontainer = getattr(position.node, "container")
-            index = cellcontainer.node_list.index(position.node)
-            index = index if position.before_node else index + 1
-            cellcontainer.insert(index, node)
-            return node
-        # if position in current SymbolTree
-        if position is not None and position.symbol_tree is not self:
-            raise RuntimeError("Position is not in current SymbolTree:", position)
-        if position is not None and position.node.get_node_type() == NodeType.Input:
+        if new_node.get_belong_symbol_tree():
+            raise ValueError(f"Node in the SymbolTree cannot be inserted into SymbolTree again: {new_node.get_name()}")
+
+        # Check if base_node in current SymbolTree
+        if base_node is not None:
+            stree = base_node.get_belong_symbol_tree()
+            if stree is not None and stree is not self:
+                raise RuntimeError(f"Position is not in current SymbolTree, node:{stree.get_ori_cls_name()}, "
+                                   f"current: {self.get_ori_cls_name()}.")
+
+        # Check if node is inserted between Input node
+        if base_node is not None and base_node.get_node_type() == NodeType.Input:
             valid = True
-            if position.before_node:
+            if before_node:
                 valid = False
-            if position.node.get_next() is not None and position.node.get_next().get_node_type() == NodeType.Input:
+            if base_node.get_next() is not None and base_node.get_next().get_node_type() == NodeType.Input:
                 valid = False
             if not valid:
-                raise RuntimeError("Can not insert a node before or between parameters:", position)
-        # unique targets, name while insert node into symbol_tree
-        node_name = self._node_name_namer.get_name(node)
-        node.set_name(node_name)
-        self._handle_custom_obj_in_normalized_args(node)
-        # _unique_targets must called after _update_args_for_unique and _update_kwargs_for_unique
-        self._unique_targets(node)
-        self._insert_node(position, node)
-        if isinstance(node, TreeNode):
-            node.symbol_tree.reg_observer(self)
-        if self._node_visitor:
-            self._node_visitor.append_node(node)
-        # update init-function-ast and construct-function-ast
-        if insert_to_ast:
-            self._insert_to_ast_while_insert_node(node, position)
-        return node
+                raise RuntimeError("Can not insert a node before or between parameters:", base_node.get_name())
+
+        # save target name, which is used to provide unique target
+        if new_node.get_targets():
+            for target in new_node.get_targets():
+                self._target_namer.add_name(str(target))
+
+        self._handle_custom_obj_in_normalized_args(new_node)
+
+        # Insert node into NodeManager
+        if node_manager is None:
+            if base_node is None:
+                raise RuntimeError("node_manager and base_node cannot both be None when inserting a node.")
+            node_manager = base_node.get_node_manager()
+
+        # set node's _belong_symbol_tree
+        new_node.set_belong_symbol_tree(self)
+
+        if node_manager is self:
+            NodeManager.insert_node(self, new_node, base_node, before_node)
+            if insert_to_ast:
+                # update init-function-ast and construct-function-ast
+                self.insert_to_ast_while_insert_node(new_node, base_node, before_node, self)
+        else:
+            node_manager.insert_node(new_node, base_node, before_node, insert_to_ast)
 
-    def append_node(self, node: Node, append_to_ast: bool = True) -> Node:
+        # register code changed event observer, which is used to update _modified flag.
+        if new_node.get_node_type() == NodeType.Tree:
+            new_node.symbol_tree.reg_observer(self)
+        elif isinstance(new_node, NodeManager):
+            new_node.reg_observer(self)
+
+        return new_node
+
+    def append_node(self, node: Node, node_manager: NodeManager = None, append_to_ast: bool = True) -> Node:
         """
         Append a node to SymbolTree.
 
@@ -680,13 +622,17 @@ class SymbolTree(Observer, Observable):
             node (Node): An instance of node to be appended.
             append_to_ast (bool): A bool indicates whether to update corresponding ast node at same time, default is
                 True.
+            node_manager (NodeManager): NodeManager those asts belong to. Default: None, means those asts belong to
+                NodeManager of symboltree's construct function.
 
         Returns:
             An instance of node which has been appended to SymbolTree.
         """
-        return self.insert_node(Position.create(self, self._tail, False), node, append_to_ast)
+        if node_manager is None:
+            node_manager = self
+        return self.insert_node(node, node_manager.get_tail(), False, node_manager, append_to_ast)
 
-    def append_origin_field(self, node: Node) -> Node:
+    def append_origin_field(self, node: Node, node_manager: NodeManager = None) -> Node:
         """
         Append an original field node to SymbolTree. An original field node represents a node created from existing
         statement in forward method, from existing ast node in ast of forward method, so ast node do not need to update
@@ -695,18 +641,16 @@ class SymbolTree(Observer, Observable):
 
         Args:
             node (Node): An instance of node to be appended.
+            node_manager (NodeManager): NodeManager those asts belong to. Default: None, means those asts belong to
+                NodeManager of symboltree's construct function.
 
         Returns:
             An instance of node which has been appended to SymbolTree.
         """
-        self._update_args_kwargs_for_unique(node)
-        if node.get_node_type() == NodeType.Output:
-            self._return = node
-        elif node.get_node_type() == NodeType.Input:
-            self._inputs.append(node)
-        return self.append_node(node, False)
+        return self.append_node(node, node_manager, False)
 
-    def append_input_node(self, ast_node, param_name: str, default: Optional[ScopedValue] = None):
+    def append_input_node(self, ast_node, param_name: str, default: Optional[ScopedValue] = None,
+                          node_manager: NodeManager = None):
         """
         Append an input node to SymbolTree corresponding to parameter of forward method of network class.
         This method is called while building SymbolTree usually.
@@ -716,13 +660,18 @@ class SymbolTree(Observer, Observable):
             param_name (str): A str represents name of parameter of forward method of network class.
             default (ScopedValue, optional): A ScopedValue represents default value of parameter. Default is None which
                 means parameter has no default value.
+            node_manager (NodeManager): NodeManager those asts belong to. Default: None, means those asts belong to
+                NodeManager of symboltree's construct function.
 
         Returns:
             An instance of input node which has been appended to SymbolTree.
         """
         if param_name == "self":
             return
-        for input_node in self._inputs:
+        # check param_name duplicated
+        if node_manager is None:
+            node_manager = self
+        for input_node in node_manager._inputs:
             targets = input_node.get_targets()
             if len(targets) != 1:
                 raise RuntimeError("targets should have 1 elements")
@@ -735,9 +684,10 @@ class SymbolTree(Observer, Observable):
             if exist_param == param_name:
                 raise RuntimeError("input duplicated:", param_name)
         input_node = Node.create_input_node(ast_node, param_name, default, name=f"input_{param_name}")
-        self.append_origin_field(input_node)
+        self.append_origin_field(input_node, node_manager)
 
-    def try_append_python_node(self, ast_scope: ast.AST, ast_node: ast.AST) -> Optional[Node]:
+    def try_append_python_node(self, ast_scope: ast.AST, ast_node: ast.AST,
+                               node_manager: NodeManager = None) -> Optional[Node]:
         """
         Try appending a python node to SymbolTree if 'ast_node' is not None and 'ast_node' is not Empty if 'ast_node' is
         a list or a dict.
@@ -746,6 +696,8 @@ class SymbolTree(Observer, Observable):
         Args:
             ast_scope (ast.AST): A ast node represents ast node of scope of node.
             ast_node (ast.AST): A ast node represents ast node.
+            node_manager (NodeManager): NodeManager those asts belong to. Default: None, means those asts belong to
+                NodeManager of symboltree's construct function.
 
         Returns:
             An instance of python node if a new node has been appended to SymbolTree else None.
@@ -754,9 +706,9 @@ class SymbolTree(Observer, Observable):
             return None
         if isinstance(ast_node, (list, dict)) and not ast_node:
             return None
-        return self.append_python_node(ast_scope, ast_node)
+        return self.append_python_node(ast_scope, ast_node, node_manager)
 
-    def append_python_node(self, ast_scope: ast.AST, ast_node: ast.AST) -> Node:
+    def append_python_node(self, ast_scope: ast.AST, ast_node: ast.AST, node_manager: NodeManager = None) -> Node:
         """
         Append a python node to SymbolTree.
         This method is called while building SymbolTree usually.
@@ -764,40 +716,50 @@ class SymbolTree(Observer, Observable):
         Args:
             ast_scope (ast.AST): A ast node represents ast node of scope of node.
             ast_node (ast.AST): A ast node represents ast node.
+            node_manager (NodeManager): NodeManager those asts belong to. Default: None, means those asts belong to
+                NodeManager of symboltree's construct function.
 
         Returns:
             An instance of python node which has been appended to SymbolTree.
         """
         logger.info("Ignoring unsupported node (%s) (%s).", type(ast_node).__name__, type(ast_scope).__name__)
-        node_name = self._node_name_namer.get_name(type(ast_node).__name__)
-        self._update_names_for_unique(ast_node)
+        node_name = type(ast_node).__name__
         node = Node.create_python_node(ast_node, node_name)
-        self._insert_node(Position.create(self, self._tail, False), node)
+        if node_manager is None or node_manager is self:
+            NodeManager.append_python_node(self, node)
+        else:
+            node_manager.append_python_node(node)
         return node
 
-    def set_output(self, return_value: str, index: int) -> Node:
+    def set_output(self, return_value: str, arg_index: int, return_idx: int = 0,
+                   node_manager: NodeManager = None) -> Node:
         """
         Update return value of return of forward method of network class.
 
         Args:
             return_value (str): A str represents new return value.
-            index (int): A int indicates which return value to be updated.
+            arg_index (int): A int indicates which value in return to be updated.
+            return_idx (int): A int indicates which return node to be updated. Default: 0.
+            node_manager (NodeManager): NodeManager those asts belong to. Default: None, means
+                symboltree's construct function.
 
         Returns:
             An instance of node represents return node after updated.
         """
-        if self._return is None:
-            raise RuntimeError("SymbolTree has no output")
-        self.set_node_arg(self._return, index, return_value)
-        return self._return
+        node_returns = NodeManager.get_returns(self) if node_manager is None else node_manager.get_returns()
+        if not node_returns:
+            raise RuntimeError("Current node_manager has no output")
+        if return_idx >= len(node_returns):
+            raise RuntimeError(f"return_idx {return_idx} should be less than return num {len(node_returns)}.")
+        node_return = node_returns[return_idx]
+        self.set_node_arg(node_return, arg_index, return_value)
+        return node_return
 
     def erase_node(self, node_or_name: Union[Node, str]) -> Node:
         """
         Erase a node from SymbolTree.
-        Note:
-            If node is depended on by other node, RuntimeError will raise.
 
-            Topological relation is updated.
+        Topological relation will be updated.
 
         Args:
             node_or_name (Union[Node, str]): An instance of node or a str represents name of node.
@@ -813,71 +775,51 @@ class SymbolTree(Observer, Observable):
         node = self._get_real_node(node_or_name)
         if node is None:
             raise RuntimeError("Node is not belong to current SymbolTree: ", node_or_name)
-        if hasattr(node, "container"):
-            cellcontainer = getattr(node, "container")
-            cellcontainer.erase(node)
-            return node
-        ret = AstModifier.erase_ast_from_function(self._root_ast, node.get_ast())
-        if not ret:
-            raise RuntimeError("node not in function ast tree.")
-        for key, value in self._nodes.items():
-            if id(value) == id(node):
-                self._nodes.pop(key)
-                value.isolate()
-                break
-        self._topo_mgr.on_erase_node(node)
+        # erase node in NodeManager
+        node_manager = node.get_node_manager()
+
+        logger.debug(f"[earse]stree: {self.get_opt_cls_name()}, "
+                     f"node_manager: {node_manager.get_manager_name()}, "
+                     f"code: {astunparse.unparse(node.get_ast()).strip()}, "
+                     f"node_name:{node.get_name()}")
+
+        if node_manager is self:
+            NodeManager.erase_node(self, node)
+            ret = AstModifier.erase_ast_from_function(self._root_ast, node.get_ast())
+            if not ret:
+                raise RuntimeError(f"erase node failed, node {node.get_name()} not in function ast tree.")
+        else:
+            node_manager.erase_node(node)
         self._deleted_node.append(node.get_name())
         return node
 
     def replace(self, old_node: Node, new_nodes: [Node]) -> Node:
         """
-        Replace an old_node with a node_tree. 'new_node' is the root node of the node_tree.
-        Note:
-            Rewrite will iterate all nodes linked to this root node and insert these nodes into symbol_tree.
-
-            Inputs of intra sub-tree nodes need to be welly set.
-
-            Inputs of inter sub-tree nodes will be updated by Rewrite automatically.
+        Replace an old_node with a node list.
 
         Args:
             old_node (Node): Node to be replaced.
-            new_nodes (list[Node]): Node tree to replace in.
+            new_nodes (list[Node]): Node list to replace in.
 
         Returns:
-            An instance of Node represents root of node_tree been replaced in.
+            Last node in new_nodes list.
 
         Raises:
             RuntimeError: If 'old_node' is isolated.
             RuntimeError: If 'old_node' is not belong to current SymbolTree.
         """
-
-        if hasattr(old_node, "container"):
-            self._replace_container_node(old_node, new_nodes)
-            return new_nodes[0]
         real_old_node = self._get_real_node(old_node)
         if real_old_node is None:
             raise RuntimeError("Old node is not belong to current SymbolTree:", old_node)
-        # get position
-        next_node: Node = old_node.get_next()
-        prev_node: Node = old_node.get_prev()
-        if prev_node is None and next_node is None:
-            raise RuntimeError("Try replacing a isolated node: ", old_node)
-        if prev_node is None:
-            position = self.before(next_node)
-        else:
-            position = self.after(prev_node)
-        # insert node first, because targets of new_node is determined after insert
-        new_tree_root = SymbolTree._link_nodes_and_find_root(new_nodes)
-        new_node = self._insert_tree(position, new_tree_root)
-        # use targets of insert tree to redirect edge
-        users = self.get_node_users(old_node)
-        if len(new_node.get_targets()) != 1:
-            raise RuntimeError("targets of new_node should have 1 elements")
-        for user in users:
-            self.set_node_arg_by_node(user[0], user[1], new_node)
-        # erase old_node after edge is redirected because node can be erased only when node is isolated topologically
+        # insert new_nodes into node_manager
+        node_manager = real_old_node.get_node_manager()
+        # insert new_nodes into NodeManager
+        base_node = old_node
+        for node in new_nodes:
+            self.insert_node(node, base_node, False, node_manager, True)
+            base_node = node
         self.erase_node(old_node)
-        return new_node
+        return new_nodes[-1]
 
     def set_node_arg(self, node: Union[Node, str], index: int, arg: Union[ScopedValue, str]):
         """
@@ -933,30 +875,234 @@ class SymbolTree(Observer, Observable):
         if out_idx >= len(targets):
             raise RuntimeError("out_idx out of range: ", out_idx)
         new_arg = targets[out_idx]
-        self.set_node_arg(real_dst_node, arg_idx, new_arg)
+        real_dst_node.set_arg(new_arg, arg_idx)
+        self._topo_mgr.on_update_arg_by_node(real_dst_node, arg_idx, real_src_node, out_idx)
+
+    def unique_name(self, name: str):
+        """Get a unique name in the symboltree"""
+        return self._target_namer.get_name(name)
+
+    def unique_func_name(self, name: str):
+        """Get a unique function name in the symboltree"""
+        if not hasattr(self._origin_network, name):
+            return name
+        suffix = 1
+        while hasattr(self._origin_network, f"{name}_{suffix}"):
+            suffix += 1
+        return f"{name}_{suffix}"
+
+    def set_node_target(self, node: Union[Node, str], index: int, target: Union[ScopedValue, str]):
+        """
+        Set target of `node` .
+
+        Args:
+            node (Union[Node, str]): Node to be modified. Can be a node or name of node.
+            index (int): Indicate which target being modified.
+            arg (Union[ScopedValue, str]): New target to been set.
 
-    def print_node_tabulate(self):
+        Raises:
+            ValueError: If `node` is not belong to current SymbolTree.
+            ValueError: If index of `node` 's target is greater than number of targets.
+        """
+
+        real_node = self._get_real_node(node)
+        if real_node is None:
+            raise ValueError("Node is not belong to current SymbolTree: ", node)
+        if isinstance(target, str):
+            target = ScopedValue.create_naming_value(target)
+        targets = node.get_targets()
+        if index >= len(targets):
+            raise ValueError(f"Index of node's target should be less than {len(targets)}, but got {index}")
+        old_target = targets[index]
+        targets[index] = target
+        node.set_targets(targets)
+        self._topo_mgr.on_update_target(node, index, old_target, target)
+
+    def all_nodes(self):
+        """
+        Get all nodes including nodes in CallFunction node, CellContainer node and sub symbol tree.
+
+        Returns:
+            A list of nodes.
+        """
+        nodes = []
+        node_managers = [self]
+        while node_managers:
+            node_manager = node_managers.pop()
+            nodes.extend(node_manager.nodes())
+            for node in node_manager.nodes():
+                if isinstance(node, NodeManager):
+                    node_managers.append(node)
+        for tree_node in self.get_tree_nodes():
+            stree = tree_node.symbol_tree
+            nodes.extend(stree.all_nodes())
+        return nodes
+
+    def get_node_from_name(self, node_name: str):
+        """
+        Get node from all NodeManagers in current symbol tree by `node_name`.
+
+        Args:
+            node_name (str): A str represents name of node as key of query.
+
+        Returns:
+            An instance of Node if found else None.
+        """
+        node_managers = [self]
+        while node_managers:
+            node_manager = node_managers.pop()
+            node = node_manager.get_node(node_name)
+            if node:
+                return node
+            for node in node_manager.nodes():
+                if isinstance(node, NodeManager):
+                    node_managers.append(node)
+        return None
+
+    def print_node_tabulate(self, all_nodes: bool = False):
+        """
+        Print nodes information and nodes' topological relations.
+
+        Args:
+            all_nodes (bool): Print nodes out of construct functions, such as nodes in CallFunction
+                nodes, CellContainer nodes and sub symbol trees.
+        """
         try:
-            from tabulate import tabulate
+            from tabulate import tabulate # pylint: disable=unused-import,reportMissingModuleSource
         except ImportError:
-            print("`print_tabular` relies on the library `tabulate`, "
-                  "which could not be found on this machine. Run `pip "
-                  "install tabulate` to install the library.")
-        node_specs = [[n.get_node_type(), n.get_name(), n.get_targets(), n.get_args(), n.get_kwargs()]
-                      for n in self.nodes()]
-        print(tabulate(node_specs, headers=['node type', 'name', 'target', 'args', 'kwargs']))
+            logger.warning("print_node_tabulate relies on the library `tabulate`, "
+                           "which could not be found on this machine. Run `pip "
+                           "install tabulate` to install the library.")
+            return ""
+        print(NodeManager.dump(self, self.get_manager_name()))
+        if all_nodes:
+            node_managers = [self]
+            while node_managers:
+                node_manager = node_managers.pop()
+                for node in node_manager.nodes():
+                    if isinstance(node, NodeManager):
+                        print(node.dump(node.get_manager_name()))
+                        node_managers.append(node)
+            for tree_node in self.get_tree_nodes():
+                stree = tree_node.symbol_tree
+                stree.print_node_tabulate(all_nodes)
 
     def dump(self):
         """Dump graph."""
         dump_st = SymbolTreeDumper(self)
         dump_st.dump()
 
-    def update_module_ast(self):
-        for node in self._external_func_ast:
-            self._module_ast.body.append(node)
-        for node in self._father_class_ast:
-            index = self._module_ast.body.index(self._class_ast)
-            self._module_ast.body.insert(index, node)
+    def check_body_exist(self, body, code_bodies):
+        """Check whether body already exist in code_bodies"""
+        # Check import ast node exist by saving import code string to self._tmp_import_strs
+        if isinstance(body, (ast.Import, ast.ImportFrom, ast.Expr)):
+            import_str = astunparse.unparse(body)
+            if import_str in self._tmp_import_strs:
+                return True
+            self._tmp_import_strs.append(import_str)
+            return False
+
+        # Check ClassDef ast node exist by using AstClassFinder
+        if isinstance(body, ast.ClassDef):
+            if sys.version_info >= (3, 9):
+                class_finder = AstClassFinder(ast.Module(body=code_bodies, type_ignores=[]))
+            else:
+                class_finder = AstClassFinder(ast.Module(body=code_bodies))
+            results = class_finder.find_all(body.name)
+            return bool(results)
+
+        # Check FunctionDef ast node exist by using AstFunctionFinder
+        if isinstance(body, ast.FunctionDef):
+            if sys.version_info >= (3, 9):
+                function_finder = AstFunctionFinder(ast.Module(body=code_bodies, type_ignores=[]))
+            else:
+                function_finder = AstFunctionFinder(ast.Module(body=code_bodies))
+            results = function_finder.find_all(body.name)
+            return bool(results)
+
+        return False
+
+    def update_class_name_of_unmodified_stree(self, stree, code_bodies) -> bool:
+        """
+        For the unmodified symbol tree, only one definition code remains in the generated code.
+        Everywhere else calling this symbol tree will use the class in this definition code.
+        """
+        # all modified ast.ClassDef will be exported to code
+        if stree.is_modified():
+            return False
+        # all un-modified ast.ClassDef only keep one instance
+        first_cls_name = self._tmp_unmodified_strees.get(type(stree.get_origin_network()))
+        if first_cls_name is None:
+            class_ast = stree.get_class_ast()
+            if class_ast:
+                self._tmp_unmodified_strees[type(stree.get_origin_network())] = class_ast.name
+            return False
+        # Un-modified ast.ClassDef already exist in code_bodies,
+        # replace class name to class name of first un-modified ast.ClassDef.
+        if sys.version_info >= (3, 9):
+            replacer = AstReplacer(ast.Module(body=code_bodies, type_ignores=[]))
+        else:
+            replacer = AstReplacer(ast.Module(body=code_bodies))
+        replacer.replace_all(stree.get_class_ast().name, first_cls_name)
+        self._tmp_replacers.append(replacer)
+        return True
+
+    def convert_stree_to_code_bodies(self, stree, code_bodies, insert_pos=0):
+        """
+        Convert nodes in stree to code_bodies
+
+        1. Add import asts into code_bodies
+        2. Add class, function and other type of asts into code_bodies
+        3. Add father class asts into code_bodies
+        4. Add external function asts into code_bodies
+        5. Add subtrees to code_bodies
+        5.1 Add subtrees in construct to code_bodies
+        5.2 Add subtrees in CellContainers to code_bodies
+
+        """
+        # Add import asts into code_bodies
+        for body in stree.get_import_asts():
+            if not self.check_body_exist(body, code_bodies):
+                code_bodies.insert(insert_pos, body)
+                insert_pos += 1
+
+        # Add class, function and other type of asts into code_bodies
+        if stree.get_module_ast():
+            for body in stree.get_module_ast().body:
+                if self.check_body_exist(body, code_bodies):
+                    continue
+                if isinstance(body, (ast.ClassDef, ast.FunctionDef)):
+                    code_bodies.insert(insert_pos, body)
+                else:
+                    code_bodies.append(body)
+
+        # Add father class asts into code_bodies
+        for body in reversed(stree.get_father_class_ast()):
+            if self.check_body_exist(body, code_bodies):
+                # remove exist ast in old position, then insert ast to upper position
+                if sys.version_info >= (3, 9):
+                    exist_ast = AstClassFinder(ast.Module(body=code_bodies, type_ignores=[])).find_all(body.name)[0]
+                else:
+                    exist_ast = AstClassFinder(ast.Module(body=code_bodies)).find_all(body.name)[0]
+                code_bodies.remove(exist_ast)
+            code_bodies.insert(insert_pos, body)
+
+        # Add external asts into code_bodies
+        for body in stree.get_external_ast():
+            if not self.check_body_exist(body, code_bodies):
+                code_bodies.insert(insert_pos, body)
+                insert_pos += 1
+
+        # Add subtrees to code_bodies
+        for node in stree.get_tree_nodes():
+            sub_stree = node.symbol_tree
+            # Ignore TreeNode create by function in the class
+            if isinstance(sub_stree.get_module_ast(), ast.FunctionDef):
+                continue
+            # For the unmodified class, update class name to name of first class
+            if self.update_class_name_of_unmodified_stree(sub_stree, code_bodies):
+                continue
+            self.convert_stree_to_code_bodies(node.symbol_tree, code_bodies, insert_pos)
 
     def get_code(self) -> str:
         """
@@ -965,34 +1111,22 @@ class SymbolTree(Observer, Observable):
         Returns:
             A str represents source code of modified network.
         """
-        self._remove_unused_import()
-        if self._init_func_ast:
-            self._remove_unused_field()
-        self._remove_duplicated_import()
-        self.update_module_ast()
+        self._tmp_import_strs.clear()
+        self._tmp_unmodified_strees.clear()
+        self._tmp_replacers.clear()
+        code_bodies = []
+        self.convert_stree_to_code_bodies(self, code_bodies)
+        if sys.version_info >= (3, 9):
+            gencode_module = ast.Module(body=code_bodies, type_ignores=[])
+        else:
+            gencode_module = ast.Module(body=code_bodies)
+        SymbolTree._remove_unused_import(gencode_module)
+        SymbolTree._remove_duplicated_import(gencode_module)
         ast.fix_missing_locations(self._module_ast)
-        # Find all ast.ClassDef which can be export to code
-        # Replace duplicated ast.ClassDef reference in main-ClassDef
-        seen_class: {type, str} = {}
-        allow_class_name = [self._class_ast.name]
-        replacers = []
-        SymbolTree._find_all_class_in_symboltree(self, seen_class, allow_class_name, replacers)
-        # Add all non-ClassDef body to gencode_module
-        # Add all ClassDef in allow_class_name to gencode_module
-        # Use gencode_module to generate code
-        bodies = []
-        for body in self._module_ast.body:
-            if not isinstance(body, ast.ClassDef):
-                bodies.append(body)
-                continue
-            if body.name in allow_class_name:
-                bodies.append(body)
-        gencode_module = ast.Module(body=bodies)
-        if_fixer = IfFixer()
-        if_fixer.fix(gencode_module)
+        IfFixer().fix(gencode_module)
         code = astunparse.unparse(gencode_module)
-        # Restore main-ClassDef
-        for replacer in replacers:
+        # Revert the class name to its original state
+        for replacer in self._tmp_replacers:
             replacer.undo_all()
         return code
 
@@ -1026,305 +1160,71 @@ class SymbolTree(Observer, Observable):
             f.write(source.encode('utf-8'))
             f.flush()
 
-    def update_scope_for_unique(self, node: Union[ast.Attribute, ast.Call, ast.Subscript]):
-        """ Update scope of ast node because of unique-ing of targets of other nodes. """
-        if isinstance(node, ast.Call):
-            self.update_scope_for_unique(node.func)
-            return
-        if not isinstance(node, (ast.Attribute, ast.Subscript)):
-            logger.warning(f"Cannot update node {astunparse.unparse(node)} for unique, type of node should "
-                           f"be one of (ast.Attribute, ast.Subscript).")
-            return
-        scope = node.value
-        if not isinstance(scope, ast.Name):
-            self.update_scope_for_unique(scope)
-            return
-        scope_name = scope.id
-        scope_name_unique = self._target_namer.get_real_arg(scope_name)
-        scope.id = scope_name_unique
-
-    def _insert_to_ast_while_insert_node(self, node: Node, position: Optional[Position]):
+    def insert_to_ast_while_insert_node(self, new_node: Node, base_node: Node, before_node: bool,
+                                        node_manager: NodeManager):
         """ insert_to_ast_while_insert_node. """
-        node.set_func(ScopedValue.create_naming_value(node.get_name(), "self"))
-        node_ast = node.get_ast()
-        if not isinstance(node_ast, ast.Assign):
-            raise RuntimeError("Only support insert cell op now")
-        if isinstance(node, TreeNode):
-            setattr(self._origin_network, node.get_name(), node.get_instance())
-            args_call = AstModifier.create_call(ScopedValue(ValueType.NamingValue, "", "getattr"),
-                                                [ScopedValue(ValueType.NamingValue, "", "obj"),
-                                                 ScopedValue(ValueType.StringValue, "", node.get_name())])
-            value = ast.Call(func=ast.Name(node.symbol_tree.get_opt_cls_name(), ast.Store(), lineno=0,
-                                           col_offset=0), args=[args_call], keywords=[], lineno=0, col_offset=0)
-
-            ast_target = ast.Name("self." + node.get_name(), ast.Store(), lineno=0, col_offset=0)
-            assign = ast.Assign(targets=[ast_target], value=value, lineno=0, col_offset=0)
-            AstModifier.insert_assign_ast_to_function(self._init_func_ast, assign)
-
-            AstModifier.insert_assign_ast_to_function(self._root_ast, node_ast,
-                                                      None if position is None else position.node.get_ast(),
-                                                      position.before_node)
-            sub_stree: SymbolTree = node.symbol_tree
-            from .symbol_tree_builder import SymbolTreeBuilder
-            SymbolTreeBuilder.merge_module_of_subtree(self, sub_stree)
+        if new_node.get_node_type() == NodeType.Input:
+            # insert a new input
+            self._inputs.append(new_node)
+            ast_construct = self.get_ast_root()
+            arg: str = new_node.get_targets()[0].value
+            ast_arg = ast.arg(arg=arg, annotation=None, type_comment=None)
+            AstModifier.append_arg_to_function(ast_construct, ast_arg)
         else:
-            AstModifier.insert_assign_to_function(self._init_func_ast,
-                                                  targets=[ScopedValue(ValueType.NamingValue, "self", node.get_name())],
-                                                  expr=ScopedValue(ValueType.NamingValue, "", "getattr"),
-                                                  args=[ScopedValue(ValueType.NamingValue, "", "obj"),
-                                                        ScopedValue(ValueType.StringValue, "", node.get_name())])
-            AstModifier.insert_assign_ast_to_function(self._root_ast, node_ast,
-                                                      None if position is None else position.node.get_ast(),
-                                                      position.before_node)
-        setattr(self._origin_network, node.get_name(), node.get_instance())
-
-    def _remove_unused_import(self):
-        """remove unused import in self._module_ast"""
-        str_checker = StrChecker(self._module_ast)
-        for i in range(len(self._module_ast.body) - 1, -1, -1):
-            body = self._module_ast.body[i]
-            if not isinstance(body, (ast.Import, ast.ImportFrom)):
-                continue
-            if isinstance(body, ast.Import):
-                continue
-            if isinstance(body, ast.ImportFrom) and body.module == "cell":
-                self._module_ast.body.remove(body)
-                continue
-            for alias in body.names:
-                name = alias.asname if alias.asname else alias.name
-                if not str_checker.check(name):
-                    if len(body.names) == 1:
-                        self._module_ast.body.remove(body)
-                        i += 1
-                    else:
-                        body.names.remove(alias)
-
-    def _replace_container_node(self, old_node, new_nodes):
-        cellcontainer = getattr(old_node, "container")
-        index = cellcontainer.node_list.index(old_node)
-        for n in reversed(new_nodes):
-            cellcontainer.insert(index, n)
-        index = cellcontainer.node_list.index(old_node)
-        cellcontainer.erase(old_node)
-
-    def _filter_out_to_delete_field(self, to_delete_field):
-        """filter out used field from `to_delete_field`"""
-        for func_def in self._class_ast.body:
-            if not isinstance(func_def, ast.FunctionDef):
-                continue
-            if func_def.name != "__init__":
-                to_delete_to_delete_keys = []
-                property_checker = CheckPropertyIsUsed(func_def)
-                for key, _ in self._deleted_field.items():
-                    if property_checker.check("self", key):
-                        to_delete_to_delete_keys.append(key)
-                        property_checker = CheckPropertyIsUsed(func_def)
-                for key in to_delete_to_delete_keys:
-                    self._deleted_field.pop(key)
+            # insert a new assign statement
+            ast_assign = new_node.get_ast()
+            if ast_assign is None:
+                func_name = new_node.get_belong_symbol_tree().unique_func_name(new_node.get_name())
+                new_node.set_func_name(ScopedValue.create_naming_value(func_name, "self"))
+                ast_assign = new_node.update_ast_node()
+            if not isinstance(ast_assign, ast.Assign):
+                raise ValueError(f"Only support insert ast.Assign or Input now, but get {type(ast_assign)}")
+            # Save instance into _origin_network.
+            setattr(self._origin_network, new_node.get_name(), new_node.get_instance())
+            # Insert ast to __init__ function
+            if isinstance(new_node, TreeNode):
+                init_code = f"self.{new_node.get_name()} = " \
+                            f"{new_node.symbol_tree.get_opt_cls_name()}(obj.{new_node.get_name()})"
             else:
-                for body in func_def.body:
-                    if not isinstance(body, ast.If):
-                        continue
-                    test = body.test
-                    field_finder = FieldFinder(test)
-                    to_delete_to_delete_keys = []
-                    for key, _ in self._deleted_field.items():
-                        if field_finder.check(key):
-                            to_delete_to_delete_keys.append(key)
-                    for key in to_delete_to_delete_keys:
-                        self._deleted_field.pop(key)
-
-    def _remove_unused_field(self):
-        """remove unused field in __init__ function"""
-        multi_targets = []
-        for index, body in enumerate(self._init_func_ast.body):
-            if not isinstance(body, ast.Assign):
-                continue
-            targets = body.targets
-            for target in targets:
-                if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name) \
-                        and target.value.id == "self":
-                    self._deleted_field[target.attr] = index
-                    if len(targets) > 1:
-                        multi_targets.append(index)
-        self._filter_out_to_delete_field(self._deleted_field)
-        for i in range(len(self._init_func_ast.body) - 1, -1, -1):
-            if i in self._deleted_field.values():
-                if i in multi_targets:
-                    raise RuntimeError("Can not erase field ast node in __init__ function because of multi-targets")
-                AstModifier.erase_ast_from_function(self._init_func_ast, self._init_func_ast.body[i])
-        ast.fix_missing_locations(self._init_func_ast)
-
-    def _remove_duplicated_import(self):
-        """Remove duplicated import of 'net'."""
-        imports = []
-        for body in self._module_ast.body:
-            if isinstance(body, (ast.ImportFrom, ast.Import)):
-                import_str = astunparse.unparse(body)
-                if import_str not in imports:
-                    imports.append(import_str)
-                else:
-                    self._module_ast.body.remove(body)
+                init_code = f"self.{new_node.get_name()} = obj.{new_node.get_name()}"
+            init_ast = ast.parse(init_code).body[0]
+            AstModifier.insert_assign_ast_to_function(self._init_func_ast, init_ast)
+            # Insert ast to construct_function/class_internal_function
+            ast_base_node = base_node.get_ast() if base_node else None
+            ast_functiondef = node_manager.get_ast_functiondef()
+            if not ast_functiondef:
+                raise RuntimeError(f"ast_functiondef is None in node_manager {node_manager.get_manager_name()} "
+                                   "when inserting the ast.")
+            AstModifier.insert_assign_ast_to_function(ast_functiondef, ast_assign, ast_base_node, before_node)
 
     def _get_real_node(self, node_or_name: Union[Node, str]) -> Optional[Node]:
         if isinstance(node_or_name, str):
             return self.get_node(node_or_name)
         return node_or_name
 
-    def _insert_tree(self, position: Position, root: Node, insert_to_ast: bool = True) -> Node:
-        """
-        Insert a node-tree into SymbolTree.
-        Note:
-            Inputs of intra sub-tree nodes need to be welly set.
-
-            Inputs of inter sub-tree nodes will be updated by Rewrite automatically.
-
-        Args:
-            position (Position): A Position indicates an insert position point.
-            root (Node): An instance of node as root of node-tree to be inserted in.
-            insert_to_ast (bool): A bool indicates whether to update corresponding ast node at same time, default is
-                True.
-
-        Returns:
-            An instance of node as root node of node-tree which has been inserted into SymbolTree.
-
-        Raises:
-            RuntimeError: If 'position' is not in current SymbolTree.
-        """
-
-        # if position not in current SymbolTree
-        if position.symbol_tree is not self:
-            raise RuntimeError("Position is not in current SymbolTree: ", position)
-
-        queue: [Node] = [root]
-        todos: [] = []
-        inputs_list: [] = []
-        while queue:
-            cur_node = queue.pop(0)
-            if cur_node in todos:
-                continue
-            todos.append(cur_node)
-            node_inputs = cur_node.get_inputs()
-            inputs_list.append(node_inputs)
-            for node_input in node_inputs:
-                if node_input is not None:
-                    queue.append(node_input)
-        todos.reverse()
-        inputs_list.reverse()
-        for index, todo in enumerate(todos):
-            self.insert_node(position, todo, insert_to_ast)
-            position = self.after(todo)
-            # relink input of node
-            original_inputs = inputs_list[index]
-            for arg_idx, original_input in enumerate(original_inputs):
-                if original_input is not None:
-                    self.set_node_arg_by_node(todo, arg_idx, original_input)
-        return root
-
-    def _unique_targets(self, node: Node):
-        """
-        Unique targets of node by _target_namer.
-
-        Args:
-            node (Node): A Node whose targets to be uniqued.
-        """
-        new_targets: [ScopedValue] = []
-        if node.get_targets() is None:
-            return
-        for target in node.get_targets():
-            if not isinstance(target, ScopedValue):
-                raise TypeError("target should be ScopedValue, got: ", type(target))
-            unique_target = self._target_namer.get_name(target.value)
-            new_targets.append(ScopedValue.create_naming_value(unique_target, target.scope))
-        node.set_targets(new_targets)
-
-    def _update_args_kwargs_for_unique(self, node: Node):
-        """
-        Update arguments and keyword arguments of node because unique-ing of targets of other nodes.
-
-        Args:
-            node (Node): A Node whose arguments and keyword arguments to be updated.
-        """
-        result: {str: ScopedValue} = {}
-        if node.get_normalized_args() is None:
-            return
-        for key, arg in node.get_normalized_args().items():
-            if not isinstance(arg, ScopedValue):
-                raise TypeError("arg should be ScopedValue, got: ", type(arg))
-            if arg.type == ValueType.NamingValue:
-                # unique name
-                new_arg = ScopedValue(arg.type, arg.scope, self._target_namer.get_real_arg(arg.value))
-                result[key] = new_arg
-            else:
-                result[key] = arg
-        node.set_normalized_args(result)
-
-    def _add_node2nodes(self, node: Node):
-        """
-        Add `node` to `_nodes` dict.
-
-        Args:
-            node (Node): A Node to be added into `_nodes`.
-
-        Raises:
-            RuntimeError: If name of the node is duplicated.
-        """
-        node_name = node.get_name()
-        if self._nodes.get(node_name) is not None:
-            raise RuntimeError("generated duplicated node name", node_name, self._nodes.get(node_name),
-                               node)
-        self._nodes[node_name] = node
-
-    def _insert_node(self, position: Optional[Position], node: Node):
-        """
-        Insert a node into SymbolTree.
-        1. Add `node` to `_nodes`.
-        2. Insert `node` to node list(source code order).
-        3. Update topological relation and update inputs of `node`.
-
-        Args:
-            position ([Position, optional]): Indicates node insert position. Position is None when inserting first node
-                of SymbolTree.
-            node (Node): A Node to be inserted into SymbolTree.
-
-        Raises:
-            RuntimeError: Position is None when _nodes of SymbolTree is not Empty. It means position can not be None
-                unless inserting first node.
-        """
-        if position is None:
-            if self._nodes:
-                raise RuntimeError("self._nodes should be empty")
-            self._head = node
-        else:
-            if position.before_node:
-                position.node.insert_before(node)
-            else:
-                position.node.insert_after(node)
-        self._tail = node
-        self._add_node2nodes(node)
-        self._topo_mgr.on_insert_node(node)
-        node.set_belong_symbol_tree(self)
-
     def _handle_custom_obj_in_normalized_args(self, node: Node):
         """
-        Convert CustomObjValue type argument to NamingValue type argument by storing custom object in global_vars dict.
+        Convert CustomObjValue type argument to NamingValue type argument by storing custom object to obj.
 
         Args:
             node (Node): A Node whose arguments and keyword arguments to be handled.
         """
-        result: {str, ScopedValue} = {}
-        for arg, value in node.get_normalized_args().items():
+        normalized_args: {str, ScopedValue} = {}
+        for key, value in node.get_normalized_args().items():
             if not isinstance(value, ScopedValue):
                 raise TypeError("value should be ScopedValue, got: ", type(value))
             if value.type == ValueType.CustomObjValue:
-                field = self._node_name_namer.get_name(f"var_{type(value.value).__name__}")
-                setattr(self._origin_network, field, value.value)
-                init_targets = [ScopedValue.create_naming_value(field, "self")]
-                AstModifier.append_global_vars_expr_to_init(self._init_func_ast, init_targets, field)
-                result[arg] = init_targets[0]
+                # Save CustomObjValue into _origin_network(i.e. obj): obj.arg_name = CustomObjValue
+                arg_name = self.unique_name(f"arg_{type(value.value).__name__}")
+                setattr(self._origin_network, arg_name, value.value)
+                # Add new code to __init__(): self.arg_name = obj.arg_name
+                new_ast = ast.parse(f"self.{arg_name} = obj.{arg_name}").body[0]
+                self._init_func_ast.body.append(new_ast)
+                # Modify node's normalized_args: CustomObjValue -> self.arg_name
+                normalized_args[key] = ScopedValue.create_naming_value(arg_name, "self")
             else:
-                result[arg] = value
-        node.set_normalized_args(result)
+                normalized_args[key] = value
+        node.set_normalized_args(normalized_args)
 
     def _get_cls_through_file(self):
         """
@@ -1336,12 +1236,14 @@ class SymbolTree(Observer, Observable):
         Returns:
             A class handle.
         """
-        self._update_container()
         file_path = os.getcwd()
         file_path = os.path.join(file_path, "rewritten_network")
         if not os.path.exists(file_path):
-            os.mkdir(file_path)
-        file_name = "{0}_{1}.py".format(self._opt_cls_name, id(self))
+            try:
+                os.mkdir(file_path, mode=0o700)
+            except FileExistsError:
+                pass
+        file_name = f"{self._opt_cls_name}_{id(self)}.py"
         network_file = os.path.join(file_path, file_name)
         with os.fdopen(os.open(network_file, os.O_WRONLY | os.O_CREAT, stat.S_IRWXU), 'wb') as f:
             source = self.get_code()
@@ -1355,15 +1257,20 @@ class SymbolTree(Observer, Observable):
 
         i = 0
         while not tmp_module:
-            try:
-                tmp_module = importlib.import_module(tmp_module_name)
-            except ModuleNotFoundError:
+            spec = importlib.util.spec_from_file_location(tmp_module_name, network_file)
+            if spec:
+                tmp_module = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(tmp_module)
+            else:
+                logger.warning(f"load module {tmp_module_name} failed, retrying.")
                 if i > 10:
                     break
-                time.sleep(0.1)
+                time.sleep(0.5)
                 i += 1
         if not tmp_module:
             logger.error(f"load module {tmp_module_name} failed.")
+        # Save new module to sys.modules to support inspect.getsource().
+        sys.modules[tmp_module_name] = tmp_module
         network_cls = getattr(tmp_module, self._opt_cls_name)
         if network_cls is None:
             raise RuntimeError("Can not find network class:", self._opt_cls_name)
@@ -1373,21 +1280,6 @@ class SymbolTree(Observer, Observable):
         self._modified = True
         self.changed(event)
 
-    def _update_container(self):
-        """Update instance of node in container."""
-        for node in self.nodes():
-            index = 0
-            if node.get_node_type() == NodeType.CellContainer:
-                for n in node.node_list:
-                    if not n.valid:
-                        continue
-                    if n.get_node_type() == NodeType.Tree:
-                        obj = n.symbol_tree.get_network()
-                        node.get_instance()[index] = obj
-                    else:
-                        node.get_instance()[index] = n.get_instance()
-                    index += 1
-
     def _cal_difference_set(self, input, other):
         """Calculate different set of two sets."""
         set1 = set(input)
@@ -1409,50 +1301,3 @@ class SymbolTree(Observer, Observable):
         primitives = self._cal_difference_set(self._origin_network._primitives.keys(), new_net._primitives.keys())
         for p in primitives:
             new_net._primitives[p] = self._origin_network._primitives[p]
-
-    def _update_names_for_unique(self, node: ast.AST):
-        """ Update names of ast nodes for unique. """
-        if isinstance(node, (ast.For, ast.If, ast.While)):
-            self._update_names_for_unique_branchs(node)
-        elif isinstance(node, ast.Assign):
-            self._update_names_for_unique(node.value)
-            for target in node.targets:
-                self._update_names_for_unique(target)
-        elif isinstance(node, ast.Call):
-            if isinstance(node.func, ast.Attribute):
-                self._update_names_for_unique(node.func.value)
-            for arg in node.args:
-                self._update_names_for_unique(arg)
-            for keyword in node.keywords:
-                self._update_names_for_unique(keyword)
-        elif isinstance(node, ast.UnaryOp):
-            self._update_names_for_unique(node.operand)
-        elif isinstance(node, ast.BinOp):
-            self._update_names_for_unique(node.left)
-            self._update_names_for_unique(node.right)
-        elif isinstance(node, (ast.Attribute, ast.Subscript, ast.Return)):
-            self._update_names_for_unique(node.value)
-        elif isinstance(node, (ast.List, ast.Tuple)):
-            for elt in node.elts:
-                self._update_names_for_unique(elt)
-        elif isinstance(node, ast.Compare):
-            for comparator in node.comparators:
-                self._update_names_for_unique(comparator)
-        elif isinstance(node, ast.Name):
-            node.id = self._target_namer.get_real_arg(node.id)
-
-    def _update_names_for_unique_branchs(self, node: Union[ast.For, ast.If, ast.While]):
-        """ Update names of ast nodes for unique with ast.For, ast.If or ast.While """
-        if isinstance(node, ast.For):
-            self._update_names_for_unique(node.target)
-            self._update_names_for_unique(node.iter)
-            for body in node.body:
-                self._update_names_for_unique(body)
-            for body in node.orelse:
-                self._update_names_for_unique(body)
-        elif isinstance(node, (ast.If, ast.While)):
-            self._update_names_for_unique(node.test)
-            for body in node.body:
-                self._update_names_for_unique(body)
-            for body in node.orelse:
-                self._update_names_for_unique(body)