diff --git a/.jenkins/check/config/filter_pylint.txt b/.jenkins/check/config/filter_pylint.txt index dde8a938513..1b1e8268c3c 100644 --- a/.jenkins/check/config/filter_pylint.txt +++ b/.jenkins/check/config/filter_pylint.txt @@ -40,6 +40,7 @@ "mindspore/mindspore/python/mindspore/log.py" "protected-access" "mindspore/mindspore/python/mindspore/rewrite/api/node.py" "protected-access" "mindspore/mindspore/python/mindspore/rewrite/node.py" "protected-access" +"mindspore/mindspore/python/mindspore/rewrite/symbol_tree.py" "protected-access" "mindspore/mindspore/python/mindspore/rewrite/parser_register.py" "protected-access" "mindspore/mindspore/python/mindspore/rewrite/api/pattern_engine.py" "protected-access" "mindspore/mindspore/python/mindspore/rewrite/symbol_tree.py" "inconsistent-return-statements" diff --git a/mindspore/python/mindspore/rewrite/__init__.py b/mindspore/python/mindspore/rewrite/__init__.py index 63bff9c65dc..e689e1b0f9b 100644 --- a/mindspore/python/mindspore/rewrite/__init__.py +++ b/mindspore/python/mindspore/rewrite/__init__.py @@ -27,6 +27,7 @@ from .api.symbol_tree import SymbolTree from .api.node import Node from .api.node_type import NodeType from .api.pattern_engine import PatternEngine, PatternNode, VarNode, Replacement +from .api.tree_node_helper import TreeNodeHelper __all__ = ["SymbolTree", "Node", "NodeType", "ScopedValue", "ValueType", "PatternEngine", "PatternNode", "VarNode", - "Replacement"] + "Replacement", "TreeNodeHelper"] diff --git a/mindspore/python/mindspore/rewrite/api/node.py b/mindspore/python/mindspore/rewrite/api/node.py index ca247a14a70..8732a2c73ff 100644 --- a/mindspore/python/mindspore/rewrite/api/node.py +++ b/mindspore/python/mindspore/rewrite/api/node.py @@ -88,8 +88,8 @@ class Node: RuntimeError: If value of kwarg in `kwargs` is not a `NamingValue`-`ScopedValue` or a `CustomObjValue`-`ScopedValue`. """ - return Node(NodeImpl.create_call_cell(cell, None, targets, ScopedValue.create_naming_value(name, "self"), args, - kwargs, name)) + return Node(NodeImpl.create_call_buildin_op(cell, None, targets, ScopedValue.create_naming_value(name, "self"), + args, kwargs, name)) def get_prev(self) -> 'Node': """ diff --git a/mindspore/python/mindspore/rewrite/api/pattern_engine.py b/mindspore/python/mindspore/rewrite/api/pattern_engine.py index b2875fef5d8..db7c94fd4be 100644 --- a/mindspore/python/mindspore/rewrite/api/pattern_engine.py +++ b/mindspore/python/mindspore/rewrite/api/pattern_engine.py @@ -359,8 +359,8 @@ class PatternEngine: """ # Don't iterate into subgraph node, pattern should not be matched across sub-tree - if node.get_node_type() != NodeType.CallCell and node.get_node_type() != NodeType.Input: - logger.debug("Pattern match failed: node(%s) is not a cell", str(node)) + if node.get_node_type() not in (NodeType.CallCell, NodeType.CallPrimitive, NodeType.Input): + logger.debug("Pattern match failed: node(%s) is not a CallCell, CallPrimitive or Input", str(node)) return False, OrderedDict() if not pattern.match(node): logger.debug("Pattern match failed: node(%s)'s type is %s while pattern type is %s", str(node), diff --git a/mindspore/python/mindspore/rewrite/api/symbol_tree.py b/mindspore/python/mindspore/rewrite/api/symbol_tree.py index fc2d3e2ca57..cbecaef05de 100644 --- a/mindspore/python/mindspore/rewrite/api/symbol_tree.py +++ b/mindspore/python/mindspore/rewrite/api/symbol_tree.py @@ -33,8 +33,14 @@ class SymbolTree: RuntimeError: If there is any unsupported ast node type while parsing or optimizing. """ - def __init__(self, network: Cell): - self._symbol_tree: SymbolTreeImpl = SymbolTreeBuilder(network).build() + def __init__(self, handler: SymbolTreeImpl): + self._symbol_tree: SymbolTreeImpl = handler + + @classmethod + def create(cls, network): + if not isinstance(network, Cell): + raise RuntimeError("Only support Cell-type-network now.") + return cls(SymbolTreeBuilder(network).build()) def get_handler(self) -> SymbolTreeImpl: """ @@ -52,7 +58,7 @@ class SymbolTree: Returns: A dict mapping from name of node to node. """ - return [Node(node_impl) for node_impl in self._symbol_tree.nodes(unfold_subtree=True)] + return [Node(node_impl) for node_impl in self._symbol_tree.nodes(unfold_subtree=False)] def get_node(self, node_name: str) -> Optional[Node]: """ diff --git a/mindspore/python/mindspore/rewrite/api/tree_node_helper.py b/mindspore/python/mindspore/rewrite/api/tree_node_helper.py new file mode 100644 index 00000000000..10a66757206 --- /dev/null +++ b/mindspore/python/mindspore/rewrite/api/tree_node_helper.py @@ -0,0 +1,55 @@ +# Copyright 2022 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Rewrite module api: TreeNodeHelper.""" +from typing import Optional + +from .symbol_tree import SymbolTree +from .node import Node +from .node_type import NodeType +from ..symbol_tree import SymbolTree as SymbolTreeImpl +from ..node import TreeNode + + +class TreeNodeHelper: + """ + `TreeNodeHelper` is used to break circle reference while getting symbol_tree from a `Tree` type `Node`. + + `TreeNodeHelper` provides a staticmethod `get_sub_tree` for getting symbol_tree from a `Tree` type `Node`. + """ + + @staticmethod + def get_sub_tree(node: Node) -> Optional[SymbolTree]: + """ + Getting symbol_tree from a `Tree` type `Node`. + + Args: + node (Node): A `Node` who may hold a sub-symbol_tree. + + Returns: + An instance of SymbolTree represents sub-symbol_tree. Note that `node`'s symbol_tree maybe None, in this + case, method will return None. + + Raises: + RuntimeError: If `node`'s type is not `NodeType.Tree`. + """ + + if node.get_node_type() == NodeType.Tree: + node_impl = node.get_handler() + assert isinstance(node_impl, TreeNode) + subtree: SymbolTreeImpl = node_impl.symbol_tree + if subtree is None: + return None + return SymbolTree(subtree) + raise RuntimeError("Current node is not a Tree node") diff --git a/mindspore/python/mindspore/rewrite/ast_helpers/ast_finder.py b/mindspore/python/mindspore/rewrite/ast_helpers/ast_finder.py index 1723b0b311c..7ea5a7df967 100644 --- a/mindspore/python/mindspore/rewrite/ast_helpers/ast_finder.py +++ b/mindspore/python/mindspore/rewrite/ast_helpers/ast_finder.py @@ -35,7 +35,6 @@ class AstFinder(ast.NodeVisitor): """ An override method, iterating over all nodes and save target ast nodes. - Args: node (ast.AST): An instance of ast node which is visited currently. """ diff --git a/mindspore/python/mindspore/rewrite/ast_helpers/ast_modifier.py b/mindspore/python/mindspore/rewrite/ast_helpers/ast_modifier.py index bbce33604a2..e940e32438c 100644 --- a/mindspore/python/mindspore/rewrite/ast_helpers/ast_modifier.py +++ b/mindspore/python/mindspore/rewrite/ast_helpers/ast_modifier.py @@ -49,12 +49,11 @@ class AstModifier(ast.NodeTransformer): Args: ast_father (ast.AST): Where new ast node to be inserted into. ast_son (ast.AST): An ast node to be inserted in. - index_ast (Optional[ast.AST]): An ast_node indicates a position in 'ast_father' where new ast node to be - inserted into. Default is None which means append new ast node to body of - 'ast_father'. + index_ast ([ast.AST, optional]): An ast_node indicates a position in 'ast_father' where new ast node to be + inserted into. Default is None which means append new ast node to body of 'ast_father'. insert_before (bool): A bool indicates at before or at after of 'index_ast' where new ast node to be - inserted into. Only valid when 'index_ast' is not None. Default is True which means - inserting new ast node before 'index_ast'. + inserted into. Only valid when 'index_ast' is not None. Default is True which means inserting new ast + node before 'index_ast'. Returns: An instance of ast.AST which has been inserted into 'ast_father'. @@ -97,12 +96,11 @@ class AstModifier(ast.NodeTransformer): expr (ScopedValue): Func of ast.Call which is value of new ast.Assign. args ([ScopedValue]): Args of ast.Call which is value of new ast.Assign. kwargs ({str, ScopedValue}): Kwargs of ast.Call which is value of new ast.Assign. - index_ast (Optional[ast.AST]): An ast_node indicates a position in 'ast_func' where new ast.Assign node to - be inserted into. Default is None which means append new ast.Assign into - 'ast_func'. + index_ast ([ast.AST, optional]): An ast_node indicates a position in 'ast_func' where new ast.Assign node to + be inserted into. Default is None which means append new ast.Assign into 'ast_func'. insert_before (bool): A bool indicates at before or at after of 'index_ast' where new ast.Assign node to be - inserted into. Only valid when 'index_ast' is not None. Default is True which means - inserting new ast.Assign before 'index_ast'. + inserted into. Only valid when 'index_ast' is not None. Default is True which means inserting new + ast.Assign before 'index_ast'. Returns: An instance of ast.Assign which has been inserted into 'ast_func'. @@ -122,12 +120,11 @@ class AstModifier(ast.NodeTransformer): Args: ast_func (ast.FunctionDef): Where new ast.Assign to be inserted into. ast_assign (ast.Assign): An instance of ast.Assign to be inserted in. - index_ast (Optional[ast.AST]): An ast_node indicates a position in 'ast_func' where new ast.Assign node to - be inserted into. Default is None which means append new ast.Assign to - 'ast_func'. + index_ast ([ast.AST, optional]): An ast_node indicates a position in 'ast_func' where new ast.Assign node to + be inserted into. Default is None which means append new ast.Assign to 'ast_func'. insert_before (bool): A bool indicates at before or at after of 'index_ast' where new ast.Assign node to be - inserted into. Only valid when 'index_ast' is not None. Default is True which means - inserting new ast.Assign before 'index_ast'. + inserted into. Only valid when 'index_ast' is not None. Default is True which means inserting new + ast.Assign before 'index_ast'. Returns: An instance of ast.Assign which has been inserted into 'ast_func'. @@ -324,17 +321,17 @@ class AstModifier(ast.NodeTransformer): Raises: TypeError: Input src_argument is not a ScopedValue RuntimeError: If 'dst_ast' is an instance of ast.Constant but type of 'src_argument' is not - ValueType.IntValue, ValueType.FloatValue or ValueType.StringValue. + ValueType.IntValue, ValueType.FloatValue or ValueType.StringValue. RuntimeError: If 'dst_ast' is an instance of ast.Name or ast.Attribute but type of 'src_argument' is not - ValueType.NamingValue. + ValueType.NamingValue. RuntimeError: When 'dst_ast' is an instance of ast.Name, scope of 'src_argument' is not empty. RuntimeError: When 'dst_ast' is an instance of ast.Attribute, value of 'dst_ast' is not an instance of - ast.Name. + ast.Name. RuntimeError: If 'dst_ast' is an instance of ast.Tuple but type of 'src_argument' is not - ValueType.TupleValue. + ValueType.TupleValue. RuntimeError: If 'dst_ast' is an instance of ast.Constant, ast.Name, ast.Attribute or ast.Tuple. RuntimeError: When 'dst_ast' is an instance of ast.Tuple, length of elts of 'dst_ast' is not equal to length - of value of 'src_argument'. + of value of 'src_argument'. """ if not isinstance(src_argument, ScopedValue): raise TypeError("src_argument should be ScopedValue, got: ", type(src_argument)) diff --git a/mindspore/python/mindspore/rewrite/node.py b/mindspore/python/mindspore/rewrite/node.py index 57698978277..204eef81c55 100644 --- a/mindspore/python/mindspore/rewrite/node.py +++ b/mindspore/python/mindspore/rewrite/node.py @@ -18,6 +18,7 @@ import ast import inspect from mindspore.nn import Cell +from mindspore.ops import Primitive from mindspore import log as logger from .ast_helpers import AstModifier from .api.scoped_value import ScopedValue, ValueType @@ -63,17 +64,17 @@ class Node: func: Optional[ScopedValue], args: [ScopedValue], kwargs: {str: ScopedValue}, name: str, instance): """ Constructor of Node. Rewrite recommend invoking class method of Node to instantiate an instance of Node such - as `create_call_cell`, `create_call_method`, `create_python_node`, `create_input_node` and `create_output_node`, - etc. rather than invoking constructor of Node directly. + as `create_call_buildin_op`, `create_call_method`, `create_python_node`, `create_input_node` and + `create_output_node`, etc. rather than invoking constructor of Node directly. Args: node_type (NodeType): A NodeType as type of Node. - ast_node (Optional[ast.AST]): An instance of ast.AST represents corresponding node in ast. `ast_node` should + ast_node (ast.AST, optional): An instance of ast.AST represents corresponding node in ast. `ast_node` should not be None except when node type is Unknown. - targets ([ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class. - func (Optional[ScopedValue]): An instance of ScopedValue. See detail in docstring of Node class. - args ([ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class. - kwargs ({str: ScopedValue}): A list of instance of ScopedValue. See detail in docstring of Node class. + 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. name (str): A string represents name of node. Name of node will be unique when inserted into SymbolTree. Name of node also used as field name in network class. instance: Object in network corresponding to this node. @@ -99,23 +100,27 @@ class Node: self._belong_tree = None @classmethod - def create_call_cell(cls, cell: Cell, ast_node: Optional[ast.AST], targets: [Union[ScopedValue, str]], - func: Union[ScopedValue, str], args: [ScopedValue] = None, kwargs: {str: ScopedValue}=None, - name: str = ""): + def create_call_buildin_op(cls, op: Union[Cell, Primitive], ast_node: Optional[ast.AST], + targets: [Union[ScopedValue, str]], func: Union[ScopedValue, str], + args: [ScopedValue] = None, kwargs: {str: ScopedValue}=None, name: str = ""): """ - Class method of Node. Instantiate an instance of node whose type is CallCell. A CallCell node represents an - invoking to cell-op. + Class method of Node. Instantiate an instance of node whose type is `CallCell` or `CallPrimitive`. + A `CallCell` node represents an invoking to cell-op. + A `CallPrimitive` node represents an invoking to primitive-op. Args: - cell (Cell): An instance of Cell corresponding to this node. - ast_node (Optional[ast.AST]): An instance of ast.AST represents corresponding node in ast. - targets ([ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class. - func (Optional[ScopedValue]): An instance of ScopedValue. See detail in docstring of Node class. - args ([ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class. - kwargs ({str: ScopedValue}): A list of instance of ScopedValue. See detail in docstring of Node class. - name (str): A string represents name of node. Name of node will be unique when inserted into SymbolTree. + op (Union[Cell, Primitive]): An instance of `Cell` or `Primitive` corresponding to this node. + 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. + name (str): A string represents name of node. Name of node will be unique when inserted into `SymbolTree`. Name of node also used as field name in network class. """ + + assert isinstance(op, (Cell, Primitive)) if args is None: args = [] if kwargs is None: @@ -127,7 +132,11 @@ class Node: new_targets = Node._handle_targets(targets) if ast_node is None: ast_node = AstModifier.create_call_assign(new_targets, func, non_custom_args, non_custom_kwargs) - return cls(NodeType.CallCell, ast_node, new_targets, func, args, kwargs, name, cell) + if isinstance(op, Cell): + node_type = NodeType.CallCell + else: + node_type = NodeType.CallPrimitive + return cls(node_type, ast_node, new_targets, func, args, kwargs, name, op) @classmethod def create_call_method(cls, ast_node: Optional[ast.AST], targets: [Union[ScopedValue, str]], @@ -138,12 +147,12 @@ class Node: invoking to cell-op. Args: - ast_node (Optional[ast.AST]): An instance of ast.AST represents corresponding node in ast. `ast_node` should - not be None currently. - targets ([ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class. - func (Optional[ScopedValue]): An instance of ScopedValue. See detail in docstring of Node class. - args ([ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class. - kwargs ({str: ScopedValue}): A list of instance of ScopedValue. See detail in docstring of Node class. + ast_node ([ast.AST, optional]): An instance of ast.AST represents corresponding node in ast. `ast_node` + should not be None currently. + 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. name (str): A string represents name of node. Name of node will be unique when inserted into SymbolTree. Name of node also used as field name in network class. """ @@ -186,7 +195,7 @@ class Node: Args: ast_node (ast.AST): An instance of ast.AST represents corresponding node in ast. arg_name (str): A string represents name of parameter. - default (Optional[ScopedValue]): An instance of ScopedValue represents default value of parameter. + default ([ScopedValue, optional]): An instance of ScopedValue represents default value of parameter. name (str): A string represents name of node. Name of node will be unique when inserted into SymbolTree. Name of node also used as field name in network class. """ @@ -205,19 +214,18 @@ class Node: Args: ast_node (ast.AST): An instance of ast.AST represents corresponding node in ast. - return_values ([str]): A list of string represents name of return values. + return_values (list[str]): A list of string represents name of return values. name (str): A string represents name of node. Name of node will be unique when inserted into SymbolTree. Name of node also used as field name in network class. """ real_return_values = ScopedValue.create_name_values(return_values) return cls(NodeType.Output, ast_node, None, ScopedValue.create_naming_value("return"), real_return_values, {}, - name, - None) + name, None) @staticmethod def _get_construct_arg_names(parameters): """ - Static method of Node. Get parameters' names of the construct function. + Static method of `Node`. Get parameters' names of the construct function. Args: parameters (MappingProxyType): An ordered mapping of parameters' names to the corresponding Parameter @@ -265,9 +273,9 @@ class Node: Args: names (tuple): Parameters' name got from construct func. - args ([ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class. - kwargs ({str: ScopedValue}): A list of instance of ScopedValue. See detail in docstring of Node class. - normalized_args ({str: ScopedValue}): The normalized args to be filled. + 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. + normalized_args (dict{str: ScopedValue}): The normalized args to be filled. Raises: RuntimeError: Input args are invalid. @@ -321,8 +329,8 @@ class Node: The keys of args are obtained from the construct function of type(self._instance). Args: - args ([ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class. - kwargs ({str: ScopedValue}): A list of 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. Raises: RuntimeError: Input args are invalid. @@ -361,7 +369,7 @@ class Node: Convert CustomObjValue type argument to NamingValue type argument. Args: - args ([ScopedValue]): A list of instance of ScopedValue to be converted. + args (list[ScopedValue]): A list of instance of ScopedValue to be converted. Returns: A list of instance of ScopedValue which have been converted. @@ -383,7 +391,7 @@ class Node: Convert CustomObjValue type argument to NamingValue type argument. Args: - kwargs ({str: ScopedValue}): A str to instance of ScopedValue dict whose value to be converted. + kwargs (dict{str: ScopedValue}): A str to instance of ScopedValue dict whose value to be converted. Returns: A str to instance of ScopedValue dict whose value has be converted. @@ -483,7 +491,7 @@ class Node: self._belong_tree = symbol_tree def _sync_assign_func_to_ast(self): - """Sync func of ast.Call of ast.Assign from self._name when NodeType is CallCell.""" + """Sync func of ast.Call of ast.Assign from self._name when NodeType is CallCell or CallPrimitive.""" if self._ast_node is None: return assign_ast = self._ast_node @@ -510,7 +518,7 @@ class Node: ast.fix_missing_locations(assign_ast) def _sync_assign_targets_to_ast(self): - """Sync targets of ast.Assign from self._targets when NodeType is CallCell or CallMethod.""" + """Sync targets of ast.Assign from self._targets when NodeType is CallCell, CallPrimitive or CallMethod.""" if self._ast_node is None: return assign_ast = self._ast_node @@ -529,7 +537,7 @@ class Node: ast.fix_missing_locations(assign_ast) def _sync_call_cell_args_to_ast(self): - """Sync args of ast.Cell of ast.Assign from self._normalized_args when NodeType is CallCell.""" + """Sync args of ast.Cell of ast.Assign from self._normalized_args when NodeType is CallCell or CallPrimitive.""" if self._ast_node is None: return assign_ast = self._ast_node @@ -680,7 +688,7 @@ class Node: Args: - inputs ([Node]): A list of instances of Node as new input nodes. + inputs (list[Node]): A list of instances of Node as new input nodes. """ self._inputs = inputs @@ -723,7 +731,7 @@ class Node: targets ([ScopedValue]): A list of instances of ScopedValue as new targets. """ self._targets = targets - if self._node_type in (NodeType.CallCell, NodeType.CallMethod): + if self._node_type in (NodeType.CallCell, NodeType.CallMethod, NodeType.CallPrimitive, NodeType.Tree): self._sync_assign_targets_to_ast() def get_func(self) -> ScopedValue: @@ -746,7 +754,7 @@ class Node: func (ScopedValue): An instance of ScopedValue as new func. """ self._func = func - if self._node_type == NodeType.CallCell: + if self._node_type in (NodeType.CallCell, NodeType.CallPrimitive): self._sync_assign_func_to_ast() def get_name(self) -> str: @@ -779,10 +787,11 @@ class Node: def get_instance_type(self) -> type: """ Get the instance_type of current node. - When node_type of current node is CallCell, instance_type is type of cell-op. - When node_type of current node is CallPrimitive, instance_type is type of primitive-op. - When node_type of current node is Tree, instance_type is type of network-cell. - When node_type of current node is Python, Input, Output or CallMethod, instance_type should be NoneType + + - When node_type of current node is CallCell, instance_type is type of cell-op. + - When node_type of current node is CallPrimitive, instance_type is type of primitive-op. + - When node_type of current node is Tree, instance_type is type of network-cell. + - When node_type of current node is Python, Input, Output or CallMethod, instance_type should be NoneType Returns: A type. @@ -792,10 +801,11 @@ class Node: def get_instance(self): """ Get the instance of current node. - When node_type of current node is CallCell, instance is an instance of Cell. - When node_type of current node is CallPrimitive, instance is an instance of primitive. - When node_type of current node is Tree, instance is an instance of network-cell. - When node_type of current node is Python, Input, Output or CallMethod, instance should be None + + - When node_type of current node is CallCell, instance is an instance of Cell. + - When node_type of current node is CallPrimitive, instance is an instance of primitive. + - When node_type of current node is Tree, instance is an instance of network-cell. + - When node_type of current node is Python, Input, Output or CallMethod, instance should be None Returns: A object. @@ -804,7 +814,7 @@ class Node: def _sync_arg(self): """Sync _normalized_args to corresponding ast node when updated.""" - if self._node_type == NodeType.CallCell: + if self._node_type in (NodeType.CallCell, NodeType.CallPrimitive, NodeType.Tree): self._sync_call_cell_args_to_ast() elif self._node_type == NodeType.Output: self._sync_return_node_to_ast() @@ -819,7 +829,7 @@ class Node: Args: arg_idx (int): Indicate which input being modified. node (Node): Node as new input. Can be a node or name of node. - out_idx (Optional[int]): Indicate which output of 'node' as new argument. Default is None which means use + out_idx ([int, optional]): Indicate which output of 'node' as new argument. Default is None which means use first output of 'node_to_link' as new input. Raises: @@ -867,7 +877,7 @@ class Node: Note that when _normalized_args is updated, corresponding ast node would be updated also. Args: - args ([ScopedValue]): New arguments to been set. + args (list[ScopedValue]): New arguments to been set. Raises: TypeError: Element of new argument is not an instance of ScopedValue. @@ -886,7 +896,7 @@ class Node: Note that when _normalized_args is updated, corresponding ast node would be updated also. Args: - kwargs ({str: ScopedValue}): New arguments to been set. + kwargs (dict{str: ScopedValue}): New arguments to been set. Raises: TypeError: Value of new argument is not an instance of ScopedValue. @@ -923,11 +933,12 @@ class Node: def get_args(self): """ Get the arguments of current node. - When node_type of current node is CallCell, CallPrimitive or Tree, arguments are corresponding to args of - ast.Call which represents arguments to invoke cell-op's forward method or primitive-op's `call()` method. - When node_type of current node is Input, arguments represents default-value of argument of function. - When node_type of current node is Output, arguments represents return values. - When node_type of current node is Python, arguments are don't-care. + + - When node_type of current node is CallCell, CallPrimitive or Tree, arguments are corresponding to args of + ast.Call which represents arguments to invoke cell-op's forward method or primitive-op's `call()` method. + - When node_type of current node is Input, arguments represents default-value of argument of function. + - When node_type of current node is Output, arguments represents return values. + - When node_type of current node is Python, arguments are don't-care. Returns: A list of instances of ScopedValue. @@ -940,9 +951,11 @@ class Node: def get_kwargs(self): """ Get the keyword arguments of current node. - When node_type of current node is CallCell, CallPrimitive or Tree, keyword arguments are corresponding to kwargs - of ast.Call which represents arguments to invoke cell-op's forward method or primitive-op's `call()` method. - When node_type of current node is Python, Input or Output, keyword arguments are don't-care. + + - When node_type of current node is CallCell, CallPrimitive or Tree, keyword arguments are corresponding to + kwargs of ast.Call which represents arguments to invoke cell-op's forward method or primitive-op's `call()` + method. + - When node_type of current node is Python, Input or Output, keyword arguments are don't-care. Returns: A dict of str to instance of ScopedValue. @@ -1052,16 +1065,16 @@ class TreeNode(Node): args: [ScopedValue], kwargs: {str: ScopedValue}, name: str, instance): """ Constructor of Node. Rewrite recommend to invoking class method of Node to instantiate an instance of Node such - as `create_call_cell`, `create_call_method`, `create_python_node`, `create_input_node` and `create_output_node`, - etc. rather than invoking constructor of Node directly. + as `create_call_buildin_op`, `create_call_method`, `create_python_node`, `create_input_node` and + `create_output_node`, etc. rather than invoking constructor of Node directly. Args: tree: An instance of SymbolTree represents a handler of sub-symbol-tree. ast_node (ast.AST): An instance of ast.AST represents corresponding node in ast. - targets ([ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class. - func (Optional[ScopedValue]): An instance of ScopedValue. See detail in docstring of Node class. - args ([ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class. - kwargs ({str: ScopedValue}): A list of instance of ScopedValue. See detail in docstring of Node class. + 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. name (str): A string represents name of node. Name of node will be unique when inserted into SymbolTree. Name of node also used as field name in network class. instance: Object in network corresponding to this node. @@ -1081,10 +1094,10 @@ class TreeNode(Node): Args: tree: An instance of SymbolTree represents a handler of sub-symbol-tree. ast_node (ast.AST): An instance of ast.AST represents corresponding node in ast. - targets ([ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class. - func (Optional[ScopedValue]): An instance of ScopedValue. See detail in docstring of Node class. - args ([ScopedValue]): A list of instance of ScopedValue. See detail in docstring of Node class. - kwargs ({str: ScopedValue}): A list of instance of ScopedValue. See detail in docstring of Node class. + 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. name (str): A string represents name of node. Name of node will be unique when inserted into SymbolTree. Name of node also used as field name in network class. instance: Object in network corresponding to this node. diff --git a/mindspore/python/mindspore/rewrite/parsers/assign_parser.py b/mindspore/python/mindspore/rewrite/parsers/assign_parser.py index a069f85a18f..4a2d2011b0f 100644 --- a/mindspore/python/mindspore/rewrite/parsers/assign_parser.py +++ b/mindspore/python/mindspore/rewrite/parsers/assign_parser.py @@ -17,17 +17,27 @@ import ast import astunparse from mindspore import log as logger +from mindspore._extends.parse.namespace import CellNamespace +from mindspore.nn import Cell +from mindspore.ops import Primitive from ..symbol_tree import SymbolTree from ..node import Node, TreeNode from ..parser import Parser from ..parser_register import reg_parser from ..api.scoped_value import ScopedValue from ..symbol_tree_builder import SymbolTreeBuilder +from ..ast_helpers import AstReplacer, AstModifier class AssignParser(Parser): """Parse ast.Assign in construct function to node of SymbolTree.""" + def __init__(self): + """Constructor""" + super(AssignParser, self).__init__() + self._cell_namespce = CellNamespace('mindspore.nn') + self._primitive_namespce = CellNamespace('mindspore.ops.operations') + def target(self): """Parse target type.""" return ast.Assign @@ -166,21 +176,98 @@ class AssignParser(Parser): results[keyword.arg] = AssignParser._create_scopedvalue(keyword.value) return results + def _is_subtree_cell(self, cell: Cell) -> bool: + assert isinstance(cell, Cell) + return not type(cell).__name__ in self._cell_namespce + @staticmethod - def _convert_ast_call_to_node(ast_node: ast.Call, father_ast_node: ast.Assign, stree: SymbolTree) -> Node: + def _find_op_and_type(func_scope, func_name, stree: SymbolTree): + """ + Get the func scope from ast.Call. + + Args: + func_scope (str): Func scope. + func_name (str): Func name. + stree (SymbolTree): Belong SymbolTree. + + Returns: + A type represents type of op and an instance represents operator instance. + """ + + if func_scope != "self": + raise NotImplementedError("Not support parse operator which is instantiated at runtime now") # todo + var_dict = stree.get_origin_network().__dict__ + for key, value in var_dict["_cells"].items(): + if key == func_name: + return type(value), value + + for key, value in var_dict["_primitives"].items(): + if key == func_name: + return type(value), value + return type(None), None + + def _update_field_in_init(self, func_scope, func_name, stree: SymbolTree, sub_tree: SymbolTree): + """ + When node is an invoking to sub-network, update value of ast.Assign of corresponding field in `__init__` method. + + Update from: + + .. code-block:: + + self.field = getattr(self._handler, "field") + + to: + + .. code-block:: + + self.field = SubNetwork(global_vars.get("field_args")) + + Args: + func_scope (str): A string represents scope of function symbol. + func_name (str): A string represents function symbol. + stree (SymbolTree): The SymbolTree corresponding to main-network. + sub_tree (SymbolTree): The SymbolTree corresponding to sub-network. + + Raises: + NotImplementedError: If `func_scope` is not "self", it means corresponding op is inited in forward method. + NotImplementedError: If targets of ast.Assign of corresponding field in `__init__` method. + """ + + if func_scope != "self": + raise NotImplementedError("Not support parse operator which is instantiated at runtime now") + init_func_ast = stree.get_init_func_ast() + class_name = sub_tree.get_opt_cls_name() + for body in init_func_ast.body: + if not isinstance(body, ast.Assign): + continue + if len(body.targets) > 1: + raise NotImplementedError("Not support multi-targets in assign now!") + target = body.targets[0] + if not isinstance(target, ast.Attribute) or not(target.value, ast.Name) or target.value.id != "self": + continue + if target.attr != func_name: + continue + global_vars_key = func_name + "_args" + stree.add_global_vars(global_vars_key, sub_tree.get_global_vars()) + args_call = AstModifier.create_call(ScopedValue.create_naming_value("get", "global_vars"), + [ScopedValue.create_variable_value(global_vars_key)]) + body.value = ast.Call(func=ast.Name(class_name, ast.Store()), args=[args_call], keywords=[]) + break + + def _convert_ast_call_to_node(self, ast_node: ast.Call, father_ast_node: ast.Assign, stree: SymbolTree) -> Node: """ Convert ast.Call to a symbol tree node. Args: - ast_node ([ast.Call]): An ast.Call of assign node in construct. - father_ast_node ([ast.Assign]): Assign node in construct. - stree ([SymbolTree]): Symbol Tree under parsing. + ast_node (ast.Call): An ast.Call of assign node in construct. + father_ast_node (ast.Assign): Assign node in construct. + stree (SymbolTree): Symbol Tree under parsing. Returns: An instance of Node in Symbol Tree. Raises: - RuntimeError: kwargs in construct function assign is unsupported. + RuntimeError: If operator instance invoked by assign is undefined. """ target = AssignParser._create_scopedvalue(father_ast_node.targets[0]) func_name = AssignParser._get_func_name(ast_node) @@ -190,19 +277,25 @@ class AssignParser(Parser): func = ScopedValue.create_naming_value(func_name, func_scope) call_args = [AssignParser._create_scopedvalue(arg) for arg in ast_node.args] call_kwargs = AssignParser._create_kwargs(ast_node.keywords) - if ast_node.keywords: - raise RuntimeError("kwargs in construct function assign is unsupported.") - obj = AssignParser._get_symbol_object(func_name, stree.get_origin_network()) - # need check if node is a callmethod, like: x = len(x) - # need check if node is a callprimitive, like: x = x * 5 - is_sub_tree = False - if is_sub_tree: - stb = SymbolTreeBuilder(obj) - new_stree = stb.build() - return TreeNode(new_stree, father_ast_node, [target], func, call_args, call_kwargs, func_name, - new_stree.get_origin_network()) - return Node.create_call_cell(obj, father_ast_node, [target], func, call_args, call_kwargs, func_name) + _, op = AssignParser._find_op_and_type(func_scope, func_name, stree) + if op is None: + raise RuntimeError("Operator instance undefined: '", ast.unparse(ast_node.func), "' of '", + ast.unparse(ast_node), "'") + if isinstance(op, Primitive): + return Node.create_call_buildin_op(op, father_ast_node, [target], func, call_args, call_kwargs, func_name) + if isinstance(op, Cell): + is_sub_tree = self._is_subtree_cell(op) + if is_sub_tree: + stb = SymbolTreeBuilder(op) + new_stree = stb.build() + self._update_field_in_init(func_scope, func_name, stree, new_stree) + replacer = AstReplacer(new_stree.get_class_ast()) + replacer.replace_all(new_stree.get_ori_cls_name(), new_stree.get_opt_cls_name()) + return TreeNode(new_stree, father_ast_node, [target], func, call_args, call_kwargs, func_name, + new_stree.get_origin_network()) + return Node.create_call_buildin_op(op, father_ast_node, [target], func, call_args, call_kwargs, func_name) + raise RuntimeError("Only support Cell operator or Primitive operator, got ", type(op).__name__) def process(self, stree: SymbolTree, node: ast.Assign): """ @@ -220,12 +313,13 @@ class AssignParser(Parser): RuntimeError: Only support one target in assign now. RuntimeError: Unsupported node type in construct function. """ + targets = node.targets if len(targets) != 1: raise RuntimeError("Only support one target in assign now") value = node.value if isinstance(value, ast.Call): - node_ = AssignParser._convert_ast_call_to_node(value, node, stree) + node_ = self._convert_ast_call_to_node(value, node, stree) stree.append_origin_field(node_) elif isinstance(value, (ast.BinOp, ast.BoolOp, ast.Subscript)): logger.warning(f"ops-call({astunparse.unparse(node)}) in assign will be supported in near feature, " diff --git a/mindspore/python/mindspore/rewrite/parsers/class_def_parser.py b/mindspore/python/mindspore/rewrite/parsers/class_def_parser.py index 93702096a8c..123ad305e58 100644 --- a/mindspore/python/mindspore/rewrite/parsers/class_def_parser.py +++ b/mindspore/python/mindspore/rewrite/parsers/class_def_parser.py @@ -16,35 +16,44 @@ import ast from mindspore import log as logger +from mindspore._extends.parse.namespace import CellNamespace from ..symbol_tree import SymbolTree from ..parser import Parser from ..parser_register import ParserRegister, reg_parser from ..api.scoped_value import ScopedValue -from ..ast_helpers import AstModifier +from ..ast_helpers import AstReplacer, AstModifier class ClassDefParser(Parser): """Parse ast.ClassDef which is subclass of Cell to SymbolTree.""" + def __init__(self): + """Constructor""" + super(ClassDefParser, self).__init__() + self._cell_namespace = CellNamespace('mindspore.nn') + def target(self): """Parse target type""" return ast.ClassDef - @staticmethod - def _process_init_func_ast(init_ast: ast.FunctionDef, ori_cls_name: str, opt_cls_name: str): + def _is_subtree_field(self, ori_net, field) -> bool: + op = getattr(ori_net, field) + assert op is not None + return not type(op).__name__ in self._cell_namespace + + def _process_init_func_ast(self, stree: SymbolTree, init_ast: ast.FunctionDef): """Process init func""" - super_index = ClassDefParser._modify_super_expr_of_init_func(init_ast, ori_cls_name, opt_cls_name) + super_index = ClassDefParser._find_super_expr_of_init_func(init_ast) ClassDefParser._modify_arguments_of_init_func(init_ast) - ClassDefParser._replace_ori_field_of_init_func(init_ast.body, super_index) + self._replace_ori_field_of_init_func(stree, init_ast.body, super_index) ClassDefParser._insert_handler_to_init_func(init_ast, super_index) @staticmethod - def _modify_super_expr_of_init_func(ast_init_fn: ast.FunctionDef, ori_cls_name: str, opt_cls_name: str) -> int: - """Modify network name in super(XXnet).__init__()""" + def _find_super_expr_of_init_func(ast_init_fn: ast.FunctionDef) -> int: + """Find index of super(XXnet).__init__() in body of init ast.FunctionDef""" if not ast_init_fn.body: return -1 super_index = -1 - super_call_args = None while True: super_index += 1 expr = ast_init_fn.body[super_index] @@ -62,14 +71,7 @@ class ClassDefParser(Parser): expr_value_func_value_func = expr_value_func_value.func if not isinstance(expr_value_func_value_func, ast.Name) or expr_value_func_value_func.id != "super": continue - super_call_args = expr_value_func_value.args break - if super_call_args is None or not isinstance(super_call_args, list) or len(super_call_args) != 2: - return super_index - super_call_arg = super_call_args[0] - if super_call_arg.id != ori_cls_name: - raise RuntimeError("super_call_arg.id should equal to ori_cls_name") - super_call_arg.id = opt_cls_name return super_index @staticmethod @@ -81,8 +83,7 @@ class ClassDefParser(Parser): kw_defaults=[], defaults=[], vararg=None, kwarg=None) ast.fix_missing_locations(ast_init_fn) - @staticmethod - def _replace_ori_field_of_init_func(bodies: [], super_index: int): + def _replace_ori_field_of_init_func(self, stree: SymbolTree, bodies: [], super_index: int): """ Replace original field in init func to self.XX = getattr(self._handler, "XX"). Only keep following two kinds of ast nodes in bodies right now: @@ -103,8 +104,8 @@ class ClassDefParser(Parser): continue # ignoring super.__init__() if isinstance(body, ast.If) and isinstance(body.test, ast.Attribute) \ and isinstance(body.test.value, ast.Name) and body.test.value.id == 'self': - ClassDefParser._replace_ori_field_of_init_func(body.body, -1) - ClassDefParser._replace_ori_field_of_init_func(body.orelse, -1) + self._replace_ori_field_of_init_func(stree, body.body, -1) + self._replace_ori_field_of_init_func(stree, body.orelse, -1) continue if not isinstance(body, ast.Assign): # if not assign node, delete body_index_to_be_deleted.append(body_index) @@ -146,15 +147,15 @@ class ClassDefParser(Parser): stree ([SymbolTree]): Symbol Tree under parsing. node ([ast.ClassDef]): An ast.ClassDef node. """ - # change class name - node.name = stree.get_opt_cls_name() + replacer = AstReplacer(node) + replacer.replace_all(stree.get_ori_cls_name(), stree.get_opt_cls_name()) stree.set_class_ast(node) for body in node.body: if isinstance(body, ast.FunctionDef): if body.name == "__init__": - ClassDefParser._process_init_func_ast(body, stree.get_ori_cls_name(), stree.get_opt_cls_name()) + self._process_init_func_ast(stree, body) stree.set_init_func_ast(body) elif body.name == "construct": parser: Parser = ParserRegister.instance().get_parser(ast.FunctionDef) diff --git a/mindspore/python/mindspore/rewrite/parsers/module_parser.py b/mindspore/python/mindspore/rewrite/parsers/module_parser.py index 3026f58fa55..f8078c61b1e 100644 --- a/mindspore/python/mindspore/rewrite/parsers/module_parser.py +++ b/mindspore/python/mindspore/rewrite/parsers/module_parser.py @@ -24,23 +24,7 @@ from mindspore import log as logger from ..symbol_tree import SymbolTree from ..parser import Parser from ..parser_register import ParserRegister, reg_parser - - -class ClassFinder(ast.NodeVisitor): - """Find all ast.ClassDef in input ast node.""" - - def __init__(self): - """Keep all found ast.ClassDef in self._classes""" - self._classes: [ast.ClassDef] = [] - - def visit_ClassDef(self, node: ast.ClassDef) -> Any: - """Iterate over all nodes and save ast.ClassDef nodes.""" - self._classes.append(node) - - def find_all_classes(self, node: ast.AST) -> [ast.ClassDef]: - """Interface of ClassFinder.""" - self.visit(node) - return self._classes +from ..ast_helpers import AstFinder class ModuleParser(Parser): @@ -53,8 +37,8 @@ class ModuleParser(Parser): @staticmethod def _find_class(ast_node: ast.Module) -> ast.ClassDef: """Find all ast.ClassDef in ast.Module, only support one ast.ClassDef in ast.Module now.""" - visitor = ClassFinder() - classes = visitor.find_all_classes(ast_node) + visitor = AstFinder(ast_node) + classes = visitor.find_all(ast.ClassDef) if not classes: raise RuntimeError("No class in module") if len(classes) > 1: diff --git a/mindspore/python/mindspore/rewrite/symbol_tree.py b/mindspore/python/mindspore/rewrite/symbol_tree.py index b35d9fbf2f9..fb600fc4599 100644 --- a/mindspore/python/mindspore/rewrite/symbol_tree.py +++ b/mindspore/python/mindspore/rewrite/symbol_tree.py @@ -24,26 +24,27 @@ from mindspore.nn import Cell from mindspore import log as logger from .node import Node, TreeNode from .api.node_type import NodeType -from .ast_helpers import AstModifier +from .ast_helpers import AstModifier, AstReplacer from .api.scoped_value import ScopedValue, ValueType from .symbol_tree_dumper import SymbolTreeDumper from .topological_manager import TopoManager -from .namer import TargetNamer, NodeNamer +from .namer import TargetNamer, NodeNamer, ClassNamer +from .common.observer import Observer class Position: - """Position indicates a source code position in one network.""" + """ + Position indicates a source code position in one network. + + Rewrite recommend using class method `create()` of position rather than constructor of Position. + + Args: + symbol_tree (SymbolTree): A handler of SymbolTree indicated position in which SymbolTree. + node (Node): A handler of Node indicated position is around which Node. + before_node (bool): A bool indicated position is before or after the 'node'. + """ def __init__(self, symbol_tree, node, before_node: bool): - """ - Constructor of Position. - Recommend to use class method of position rather than constructor of Position. - - Args: - symbol_tree (SymbolTree): A handler of SymbolTree indicated position in which SymbolTree. - node (Node): A handler of Node indicated position is around which Node. - before_node (bool): A bool indicated position is before or after the 'node'. - """ self.symbol_tree = symbol_tree self.node = node self.before_node = before_node @@ -66,18 +67,20 @@ class Position: return Position(symbol_tree, node, before_node) -class SymbolTree: - """A symbol-tree usually corresponding to forward method of a network.""" +class SymbolTree(Observer): + """ + A symbol-tree usually corresponding to forward method of a network. + + Rewrite recommend using SymbolTreeBuilder to instantiate an instance of SymbolTree rather than invoking constructor + of SymbolTree directly. + + Args: + origin_network (Cell): A handler to original network instance. + module_ast (ast.Module): An instance of ast.AST represents ast node of original network. + """ def __init__(self, origin_network: Cell, module_ast: ast.Module): - """ - Constructor of SymbolTree. Rewrite recommend using SymbolTreeBuilder to instantiate an instance of SymbolTree - rather than invoking constructor of SymbolTree directly. - - Args: - origin_network (Cell): A handler to original network instance. - module_ast (ast.Module): An instance of ast.AST represents ast node of original network. - """ + super().__init__() origin_network_key = "handler" # init unique-namers self._target_namer = TargetNamer() @@ -86,13 +89,14 @@ class SymbolTree: # _node_name_namer. self._node_name_namer.add_name(origin_network_key) self._topo_mgr = TopoManager() + self._topo_mgr.reg_observer(self) self._global_vars: {str, object} = {origin_network_key: origin_network} self._nodes: {str, Node} = {} # parameters of forward method self._inputs: [Node] = [] self._ori_cls_name = type(origin_network).__name__ - self._opt_cls_name = self._ori_cls_name + "Opt" + self._opt_cls_name = ClassNamer.instance().get_name(self._ori_cls_name) self._origin_network = origin_network self._module_ast: ast.Module = module_ast self._class_ast: Optional[ast.ClassDef] = None @@ -105,6 +109,11 @@ class SymbolTree: self._tail = None self._return: Optional[Node] = None + self._modified = False + + def _on_change(self): + self._modified = True + def get_ori_cls_name(self) -> str: """ Get class name of original network. @@ -132,6 +141,16 @@ class SymbolTree: """ return self._module_ast + def set_module_ast(self, ast_node: ast.Module): + """ + Setter of _module_ast. + + Args: + ast_node (ast.Module): An instance of ast.Module represents ast node of module of corresponding network + class. + """ + self._module_ast = ast_node + def get_ast_root(self): """ Getter of `_root_ast`. @@ -141,28 +160,9 @@ class SymbolTree: """ return self._root_ast - def set_class_ast(self, ast_node: ast.ClassDef): - """ - Setter of `_class_ast`. - - Args: - ast_node (ast.ClassDef): An instance of ast.ClassDef represents ast node of corresponding network class. - """ - self._class_ast = ast_node - - def set_init_func_ast(self, ast_node: ast.FunctionDef): - """ - Setter of `_init_func_ast`. - - Args: - ast_node (ast.FunctionDef): An instance of ast.FunctionDef represents ast node of init method of - corresponding network class. - """ - self._init_func_ast = ast_node - def set_ast_root(self, ast_node: ast.FunctionDef): """ - Setter of `_root_ast`. + Setter of _root_ast. Args: ast_node (ast.FunctionDef): An instance of ast.FunctionDef represents ast node of forward method of @@ -170,6 +170,43 @@ class SymbolTree: """ self._root_ast = ast_node + def get_class_ast(self): + """ + Getter of `_class_ast`. + + Returns: + An instance of ast.ClassDef represents ast node of corresponding network class. + """ + return self._class_ast + + def set_class_ast(self, ast_node: ast.ClassDef): + """ + Setter of `_init_func_ast`. + + Args: + ast_node (ast.ClassDef): An instance of ast.ClassDef represents ast node of corresponding network class. + """ + self._class_ast = ast_node + + def get_init_func_ast(self): + """ + Getter of _init_func_ast. + + Returns: + An instance of ast.FunctionDef represents ast node of init method of corresponding network class. + """ + return self._init_func_ast + + def set_init_func_ast(self, ast_node: ast.FunctionDef): + """ + Setter of _init_func_ast. + + Args: + ast_node (ast.FunctionDef): An instance of ast.FunctionDef represents ast node of init method of + corresponding network class. + """ + self._init_func_ast = ast_node + def get_inputs(self): """ Getter of `_inputs` which represents parameters of current forward method. @@ -206,7 +243,15 @@ class SymbolTree: """ return self._origin_network - def nodes(self, unfold_subtree=True): + def get_global_vars(self): + return self._global_vars + + def add_global_vars(self, key: str, value): + if self._global_vars.get(key) is not None: + raise RuntimeError("Key of global_vars duplicated:", key) + self._global_vars[key] = value + + def nodes(self, unfold_subtree=False): """ Getter of nodes if current SymbolTree. @@ -220,7 +265,7 @@ class SymbolTree: nodes = [] for _, v in self._nodes.items(): if isinstance(v, TreeNode): - nodes.extend(self.nodes(v.symbol_tree)) + nodes.extend(v.symbol_tree.nodes()) else: nodes.append(v) return nodes @@ -241,7 +286,8 @@ class SymbolTree: def _get_real_node(self, node_or_name: Union[Node, str]) -> Optional[Node]: if isinstance(node_or_name, Node): - return self.get_node(node_or_name.get_name()) + result = self.get_node(node_or_name.get_name()) + return result if result is node_or_name else None if isinstance(node_or_name, str): return self.get_node(node_or_name) return None @@ -435,7 +481,7 @@ class SymbolTree: Args: param_name (str): A str represents name of parameter of forward method of network class. - default (Optional[ScopedValue] ): A ScopedValue represents default value of parameter. Default is None which + default (ScopedValue, optional): A ScopedValue represents default value of parameter. Default is None which means parameter has no default value. Returns: @@ -556,7 +602,7 @@ class SymbolTree: 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. + True. Returns: An instance of node as root node of node-tree which has been inserted into SymbolTree. @@ -598,10 +644,11 @@ class SymbolTree: 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 ([Node]): A list of instance of Node created by Replacement. + nodes (list[Node]): A list of instance of Node created by Replacement. Returns: An instance of Node represents root of input nodes. @@ -653,7 +700,7 @@ class SymbolTree: Args: old_node (Node): Node to be replaced. - new_nodes ([Node]): Node tree to replace in. + new_nodes (list[Node]): Node tree to replace in. Returns: An instance of Node represents root of node_tree been replaced in. @@ -717,7 +764,7 @@ class SymbolTree: dst_node (Node): Node to be modified. Can be a node or name of node. arg_idx (int): Indicate which input being modified. src_node (Node): Node as new input. Can be a node or name of node. - out_idx (Optional[int]): Indicate which output of 'src_node' as new input of 'dst_node'. Default is None + out_idx ([int, optional]): Indicate which output of 'src_node' as new input of 'dst_node'. Default is None which means use first output of 'node_to_link' as new input. Raises: @@ -757,7 +804,42 @@ class SymbolTree: A str represents source code of modified network. """ ast.fix_missing_locations(self._module_ast) - return astunparse.unparse(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 = [] + replacer = AstReplacer(self._class_ast) + for node in self.nodes(): + if not isinstance(node, TreeNode): + continue + sub_stree: SymbolTree = node.symbol_tree + # all modified ast.ClassDef should export to code + if sub_stree._modified: + allow_class_name.append(sub_stree._class_ast.name) + 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) + allow_class_name.append(self._class_ast.name) + # 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) + code = astunparse.unparse(gencode_module) + # Restore main-ClassDef + replacer.undo_all() + return code def get_network(self): """ @@ -831,8 +913,8 @@ class SymbolTree: 3. Update topological relation and update inputs of `node`. Args: - position (Optional[Position]): Indicates node insert position. Position is None when inserting first node of - SymbolTree. + 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: diff --git a/mindspore/python/mindspore/rewrite/symbol_tree_builder.py b/mindspore/python/mindspore/rewrite/symbol_tree_builder.py index 6fd65889dc0..6b3b3141559 100644 --- a/mindspore/python/mindspore/rewrite/symbol_tree_builder.py +++ b/mindspore/python/mindspore/rewrite/symbol_tree_builder.py @@ -13,27 +13,30 @@ # limitations under the License. # ============================================================================ """SymbolTree builder.""" +from copy import copy from typing import Optional import ast import inspect from mindspore.nn import Cell from .symbol_tree import SymbolTree +from .node import TreeNode from .parser_register import ParserRegister from .parser import Parser from .ast_transformers import FlattenRecursiveStmt +from .ast_helpers import AstModifier +from .ast_helpers import AstFinder class SymbolTreeBuilder: - """SymbolTree builder.""" + """ + `SymbolTreeBuilder` for building a SymbolTree from network. + + Args: + network (Cell): An instance of Cell represents a network from which SymbolTree will be built. + """ def __init__(self, network: Cell): - """ - Constructor of SymbolTreeBuilder. - - Args: - network (Cell): An instance of Cell represents a network from which SymbolTree will be built. - """ if not isinstance(network, Cell): raise RuntimeError("Only support network with Cell type now") self._origin_net = network @@ -57,6 +60,114 @@ class SymbolTreeBuilder: ast_root = transformer.transform(ast_root) return ast_root + @staticmethod + def _merge_import_of_module(main_mod: ast.Module, sub_mod: ast.Module): + """ + Merge imports of ast.Module of sub-network to ast.Module of main-network. + + Note: + Imports of sub_module would be added ahead of imports in main_module. + + Error will occur if import name is over-load because of alise. + + Args: + main_mod (ast.Module): An ast.Module corresponding to main-network. + sub_mod (ast.Module): An ast.Module corresponding to sub-network. + + """ + + sub_mod_finder = AstFinder(sub_mod) + main_mod_finder = AstFinder(main_mod) + imports_in_sub = copy(sub_mod_finder.find_all((ast.Import, ast.ImportFrom))) + imports_in_main = copy(main_mod_finder.find_all((ast.Import, ast.ImportFrom))) + assert imports_in_main + first_import = imports_in_main[0] + for clazz in imports_in_sub: + AstModifier.insert_sub_ast(main_mod, clazz, first_import, True) + + @staticmethod + def _merge_class_of_module(main_mod: ast.Module, sub_mod: ast.Module): + """ + Merge classes of ast.Module of sub-network to ast.Module of main-network. + + Note: + Classes of sub_module would be added ahead of classes in main_module. + + Args: + main_mod (ast.Module): An ast.Module corresponding to main-network. + sub_mod (ast.Module): An ast.Module corresponding to sub-network. + + """ + + sub_mod_finder = AstFinder(sub_mod) + main_mod_finder = AstFinder(main_mod) + classes_in_sub = copy(sub_mod_finder.find_all(ast.ClassDef)) + classes_in_main = copy(main_mod_finder.find_all(ast.ClassDef)) + assert classes_in_main + first_class = classes_in_main[0] + for clazz in classes_in_sub: + AstModifier.insert_class_into_module(main_mod, clazz, first_class, True) + + def _merge_module_of_subtree(self): + """ + Merge ast.Module of all sub-networks into ast.Module of main-network. + + 1. Merge imports of ast.Module. + 2. Merge classes of ast.Module. + 3. Use merged ast.Module as module of main-network and sub-network. + """ + + father_mod = self._root_tree.get_module_ast() + for node in self._root_tree.nodes(): + if isinstance(node, TreeNode): + sub_stree: SymbolTree = node.symbol_tree + SymbolTreeBuilder._merge_import_of_module(father_mod, sub_stree.get_module_ast()) + SymbolTreeBuilder._merge_class_of_module(father_mod, sub_stree.get_module_ast()) + sub_stree.set_module_ast(father_mod) + + def _reduce_redundant_import(self): + """ + Reduce redundant imports of ast.Module. + + Redundant imports may be introduced into ast.Module while merging ast.Module of sub-network to main-network. + """ + + module: ast.Module = self._root_tree.get_module_ast() + import_list = [] + exist_import = [] + exist_import_from = [] + for body in module.body: + if isinstance(body, ast.Import): + names = body.names + for name in names: + assert isinstance(name, ast.alias) + import_hash = hash((name.name, name.asname)) + if import_hash in exist_import: + continue + exist_import.append(import_hash) + import_list.append(ast.Import(names=[ast.alias(name=name.name, asname=name.asname)])) + if isinstance(body, ast.ImportFrom): + import_module = body.module + names = body.names + for name in names: + assert isinstance(name, ast.alias) + import_hash = hash((import_module, name.name, name.asname)) + if import_hash in exist_import_from: + continue + exist_import_from.append(import_hash) + import_list.append(ast.ImportFrom(module=import_module, + names=[ast.alias(name=name.name, asname=name.asname)], + level=0)) + insert_pos = None + for i in range(len(module.body) - 1, -1, -1): + body = module.body[i] + if not isinstance(body, (ast.Import, ast.ImportFrom)): + insert_pos = body + continue + module.body.pop(i) + for import_ast in import_list: + AstModifier.insert_sub_ast(module, import_ast, insert_pos, True) + def build(self) -> SymbolTree: """ Build SymbolTree. @@ -64,10 +175,15 @@ class SymbolTreeBuilder: Returns: An instance of SymbolTree. """ + self._ast_root = SymbolTreeBuilder._ast_transform(self._ast_root) if not isinstance(self._ast_root, ast.Module): raise RuntimeError("ast_root should be a ast.Module") self._root_tree: SymbolTree = SymbolTree(self._origin_net, self._ast_root) parser: Parser = ParserRegister.instance().get_parser(ast.Module) parser.process(self._root_tree, self._ast_root) + self._merge_module_of_subtree() + self._reduce_redundant_import() + ast.fix_missing_locations(self._root_tree.get_module_ast()) + self._root_tree.start_observe() return self._root_tree diff --git a/mindspore/python/mindspore/rewrite/topological_manager.py b/mindspore/python/mindspore/rewrite/topological_manager.py index 574e732cb4a..2486a8cfd45 100644 --- a/mindspore/python/mindspore/rewrite/topological_manager.py +++ b/mindspore/python/mindspore/rewrite/topological_manager.py @@ -17,9 +17,10 @@ from typing import Tuple from .api.scoped_value import ScopedValue from .node import Node +from .common.observable import Observable -class TopoManager: +class TopoManager(Observable): """SymbolTree topological-relationship manager.""" def __init__(self): @@ -30,6 +31,7 @@ class TopoManager: Value of dict is a tuple whose first is an instance of Node, whose second is an index. It means node's index th arg is argument """ + super().__init__() self._target_provider: {ScopedValue, (Node, int)} = {} self._target_consumer: {ScopedValue, [(Node, int)]} = {} @@ -137,6 +139,7 @@ class TopoManager: for index, arg in enumerate(node.get_normalized_args().values()): self._add_consumer(arg, node, index) self._update_node_inputs(node) + self.changed() def on_erase_node(self, node: Node): """ @@ -156,6 +159,7 @@ class TopoManager: self._erase_consumer(arg, node) # clear inputs of node rather than call _update_node_inputs because node is already erase from consumer dict node.set_inputs([]) + self.changed() def on_update_arg(self, node: Node, arg_idx: int, old_arg: ScopedValue, new_arg: ScopedValue): """ @@ -171,6 +175,7 @@ class TopoManager: self._erase_consumer(old_arg, node) self._add_consumer(new_arg, node, arg_idx) self._update_node_inputs(node) + self.changed() def dump(self, title=""): """ diff --git a/tests/ut/python/rewrite/test_net_simple.py b/tests/ut/python/rewrite/test_net_simple.py index ebebc8eb9c7..8a42f65a7f3 100644 --- a/tests/ut/python/rewrite/test_net_simple.py +++ b/tests/ut/python/rewrite/test_net_simple.py @@ -132,10 +132,9 @@ def test_simple_net(): Expectation: Result of rewrite can be compiled. """ net = SimpleNet(10) - stree = SymbolTree(net) + stree = SymbolTree.create(net) transform(stree) - print("------------------------------------ keys of global_vars: ", - getattr(stree.get_handler(), "_global_vars").keys()) + print("------------------------------------ keys of global_vars: ", stree.get_handler().get_global_vars().keys()) net_opt = stree.get_network() data_in = Tensor(np.ones([1, 1, 32, 32]), mindspore.float32) _cell_graph_executor.compile(net_opt, data_in) diff --git a/tests/ut/python/rewrite/test_node.py b/tests/ut/python/rewrite/test_node.py index 035d08d0f8e..52815435059 100644 --- a/tests/ut/python/rewrite/test_node.py +++ b/tests/ut/python/rewrite/test_node.py @@ -36,13 +36,13 @@ class FakeCell3(Cell): def test_create_by_cell(): """ - Feature: Python api create_call_cell of Node of Rewrite. - Description: Call create_call_cell to create a node. + Feature: Python api create_call_buildin_op of Node of Rewrite. + Description: Call create_call_buildin_op to create a CallCell node. Expectation: Success. """ - node = Node.create_call_cell(FakeCell(), None, ['x'], 'new_conv', - [ScopedValue.create_naming_value('x'), ScopedValue.create_variable_value(1)], - {"cool_boy": ScopedValue.create_naming_value('Naroto')}, 'new_conv') + node = Node.create_call_buildin_op(FakeCell(), None, ['x'], 'new_conv', + [ScopedValue.create_naming_value('x'), ScopedValue.create_variable_value(1)], + {"cool_boy": ScopedValue.create_naming_value('Naroto')}, 'new_conv') assert node._args_num == 2 assert node._kwargs_num == 1 assert node._normalized_args_keys == ["input1", "input2", "cool_boy"] @@ -84,15 +84,15 @@ def test_create_by_cell(): def test_create_by_cell2(): """ - Feature: Python api create_call_cell of Node of Rewrite. - Description: Call create_call_cell to create a node. + Feature: Python api create_call_buildin_op of Node of Rewrite. + Description: Call create_call_buildin_op to create a CallCell node. Expectation: Success. """ - node = Node.create_call_cell(FakeCell2(), None, ['x'], 'new_conv', - [ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x"), - ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x"), - ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x")], - {"cool_boy": ScopedValue.create_naming_value('Naroto')}, 'new_conv') + node = Node.create_call_buildin_op(FakeCell2(), None, ['x'], 'new_conv', + [ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x"), + ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x"), + ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x")], + {"cool_boy": ScopedValue.create_naming_value('Naroto')}, 'new_conv') assert node.get_normalized_args() == { "a": ScopedValue.create_naming_value('x'), "b": ScopedValue.create_naming_value('x'), @@ -106,15 +106,16 @@ def test_create_by_cell2(): def test_create_by_cell3(): """ - Feature: Python api create_call_cell of Node of Rewrite. - Description: Call create_call_cell to create a node. + Feature: Python api create_call_buildin_op of Node of Rewrite. + Description: Call create_call_buildin_op to create a CallCell node. Expectation: Success. """ - node = Node.create_call_cell(FakeCell3(), None, ['x'], 'new_conv', - [ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x"), - ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x")], - {"h": ScopedValue.create_naming_value(1), "f": ScopedValue.create_naming_value(2), - "cool_boy": ScopedValue.create_naming_value('Naroto')}, 'new_conv') + node = Node.create_call_buildin_op(FakeCell3(), None, ['x'], 'new_conv', + [ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x"), + ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x")], + {"h": ScopedValue.create_naming_value(1), + "f": ScopedValue.create_naming_value(2), + "cool_boy": ScopedValue.create_naming_value('Naroto')}, 'new_conv') assert node.get_normalized_args() == { "a": ScopedValue.create_naming_value('x'), "b": ScopedValue.create_naming_value('x'), diff --git a/tests/ut/python/rewrite/test_pattern_engine.py b/tests/ut/python/rewrite/test_pattern_engine.py index 2c7ebaac219..73630d12208 100644 --- a/tests/ut/python/rewrite/test_pattern_engine.py +++ b/tests/ut/python/rewrite/test_pattern_engine.py @@ -83,7 +83,7 @@ def test_one_to_one_pattern(): super().__init__([BatchNorm2d], BnReplacement()) net = ChainNetwork() - stree = SymbolTree(net) + stree = SymbolTree.create(net) conv = stree.get_node("conv") bn = stree.get_node("bn") relu1 = stree.get_node("relu1") @@ -158,7 +158,7 @@ def test_one_to_multi_chain_pattern(): super().__init__([BatchNorm2d], BnReplacement()) net = ChainNetwork() - stree = SymbolTree(net) + stree = SymbolTree.create(net) conv = stree.get_node("conv") bn = stree.get_node("bn") relu1 = stree.get_node("relu1") @@ -281,7 +281,7 @@ def test_tree_pattern(): super().__init__([Add, ReLU], AddReluReplacement()) net = TreeNetwork() - stree = SymbolTree(net) + stree = SymbolTree.create(net) conv1 = stree.get_node("conv1") conv2 = stree.get_node("conv2") add = stree.get_node("add") @@ -468,7 +468,7 @@ def test_multi_input_to_multi_pattern_tree_pattern(): """ net = TreeNetwork2() - stree = SymbolTree(net) + stree = SymbolTree.create(net) conv1 = stree.get_node("conv1") conv2 = stree.get_node("conv2") add1 = stree.get_node("add1") @@ -585,7 +585,7 @@ def test_one_input_to_multi_pattern_tree_pattern(): """ net = TreeNetwork3() - stree = SymbolTree(net) + stree = SymbolTree.create(net) conv1 = stree.get_node("conv1") conv2 = stree.get_node("conv2") add1 = stree.get_node("add1") diff --git a/tests/ut/python/rewrite/test_subtree_net.py b/tests/ut/python/rewrite/test_subtree_net.py new file mode 100644 index 00000000000..98318afa050 --- /dev/null +++ b/tests/ut/python/rewrite/test_subtree_net.py @@ -0,0 +1,123 @@ +# Copyright 2022 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +import numpy as np + +import mindspore +from mindspore import Tensor, nn +from mindspore.ops import operations as P +from mindspore.rewrite import SymbolTree, ScopedValue, Node, NodeType, TreeNodeHelper +from mindspore.common.api import _cell_graph_executor + + +class SubNet(nn.Cell): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(1, 10, 3) + self.bn = nn.BatchNorm2d(10) + + def construct(self, x): + x = self.conv(x) + x = self.bn(x) + return x + + +class MainNet(nn.Cell): + def __init__(self): + super(MainNet, self).__init__() + self.conv1 = SubNet() + self.conv2 = SubNet() + self.add = P.Add() + + def construct(self, x): + x1 = self.conv1(x) + x2 = self.conv2(x) + x = self.add(x1, x2) + return x + + +def add_relu_in_conv1(stree: SymbolTree): + for node in stree.nodes(): + if node.get_node_type() != NodeType.Tree: + continue + if node.get_name() == "conv1": + modify_stree: SymbolTree = TreeNodeHelper.get_sub_tree(node) + for inner_node in modify_stree.nodes(): + if inner_node.get_node_type() != NodeType.Output: + continue + position = modify_stree.before(inner_node) + new_relu = nn.ReLU() + new_relu_node = Node.create_call_cell(new_relu, targets=['x'], name='new_relu', + args=[ScopedValue.create_naming_value('x')]) + modify_stree.insert(position, new_relu_node) + modify_stree.set_output(0, new_relu_node.get_targets()[0].value) + break + break + + +def replace_bn_in_conv2(stree: SymbolTree): + for node in stree.nodes(): + if node.get_node_type() != NodeType.Tree: + continue + if node.get_name() == "conv2": + modify_stree: SymbolTree = TreeNodeHelper.get_sub_tree(node) + for inner_node in modify_stree.nodes(): + if inner_node.get_instance_type() != nn.BatchNorm2d: + continue + new_relu = nn.ReLU() + new_relu_node = Node.create_call_cell(new_relu, targets=['x'], name='new_relu', + args=inner_node.get_args(), kwargs=inner_node.get_kwargs()) + modify_stree.replace(inner_node, [new_relu_node]) + break + break + + +def erase_relu_in_conv2(stree: SymbolTree): + for node in stree.nodes(): + if node.get_node_type() != NodeType.Tree: + continue + if node.get_name() == "conv2": + modify_stree: SymbolTree = TreeNodeHelper.get_sub_tree(node) + for inner_node in modify_stree.nodes(): + if inner_node.get_instance_type() != nn.ReLU: + continue + assert len(inner_node.get_args()) == 1 + arg = inner_node.get_args()[0] + modify_stree.set_output(0, arg.value) + modify_stree.erase_node(inner_node) + break + break + + +def transform(stree: SymbolTree): + add_relu_in_conv1(stree) + replace_bn_in_conv2(stree) + erase_relu_in_conv2(stree) + + +def test_subtree_net(): + """ + Feature: Rewrite package api: sub-tree. + Description: Use Rewrite to parse and transform a network with sub-network. + Expectation: Rewrite can parse a network with sub-network and can modify node in sub-network successfully. + """ + + net = MainNet() + stree = SymbolTree.create(net) + transform(stree) + print(stree.get_code()) + print(stree.get_handler().get_global_vars().keys()) + net_opt = stree.get_network() + data_in = Tensor(np.ones([1, 1, 32, 32]), mindspore.float32) + _cell_graph_executor.compile(net_opt, data_in) diff --git a/tests/ut/python/rewrite/test_symbol_tree.py b/tests/ut/python/rewrite/test_symbol_tree.py index edc87d614d6..6599dd78804 100644 --- a/tests/ut/python/rewrite/test_symbol_tree.py +++ b/tests/ut/python/rewrite/test_symbol_tree.py @@ -64,31 +64,31 @@ def create_symbol_tree(): stree.set_init_func_ast(ast_init_func) stree.set_ast_root(ast_construct_func) stree.append_input_node("x") - conv_node = Node.create_call_cell(net.conv, ast_conv, [ScopedValue.create_naming_value("x")], - ScopedValue.create_naming_value("conv", "self"), - [ScopedValue.create_naming_value("x")], - {}, - "conv") + conv_node = Node.create_call_buildin_op(net.conv, ast_conv, [ScopedValue.create_naming_value("x")], + ScopedValue.create_naming_value("conv", "self"), + [ScopedValue.create_naming_value("x")], + {}, + "conv") stree.append_origin_field(conv_node) - bn_node = Node.create_call_cell(net.bn, ast_bn, [ScopedValue.create_naming_value("x")], - ScopedValue.create_naming_value("bn", "self"), - [ScopedValue.create_naming_value("x")], {}, - "bn") + bn_node = Node.create_call_buildin_op(net.bn, ast_bn, [ScopedValue.create_naming_value("x")], + ScopedValue.create_naming_value("bn", "self"), + [ScopedValue.create_naming_value("x")], {}, + "bn") bn_node = stree.append_origin_field(bn_node) - relu1_node = Node.create_call_cell(net.relu1, ast_relu1, [ScopedValue.create_naming_value("x")], - ScopedValue.create_naming_value("relu1", "self"), - [ScopedValue.create_naming_value("x")], - {}, "relu1") + relu1_node = Node.create_call_buildin_op(net.relu1, ast_relu1, [ScopedValue.create_naming_value("x")], + ScopedValue.create_naming_value("relu1", "self"), + [ScopedValue.create_naming_value("x")], + {}, "relu1") relu1_node = stree.append_origin_field(relu1_node) - relu2_node = Node.create_call_cell(net.relu2, ast_relu2, [ScopedValue.create_naming_value("x")], - ScopedValue.create_naming_value("relu2", "self"), - [ScopedValue.create_naming_value("x")], - {}, "relu2") + relu2_node = Node.create_call_buildin_op(net.relu2, ast_relu2, [ScopedValue.create_naming_value("x")], + ScopedValue.create_naming_value("relu2", "self"), + [ScopedValue.create_naming_value("x")], + {}, "relu2") relu2_node = stree.append_origin_field(relu2_node) - relu3_node = Node.create_call_cell(net.relu3, ast_relu3, [ScopedValue.create_naming_value("x")], - ScopedValue.create_naming_value("relu3", "self"), - [ScopedValue.create_naming_value("x")], - {}, "relu3") + relu3_node = Node.create_call_buildin_op(net.relu3, ast_relu3, [ScopedValue.create_naming_value("x")], + ScopedValue.create_naming_value("relu3", "self"), + [ScopedValue.create_naming_value("x")], + {}, "relu3") stree.append_origin_field(relu3_node) node_return = Node.create_output_node(ast_return, ["x"]) stree.append_origin_field(node_return) @@ -113,9 +113,10 @@ def test_insert_node(): assert len(relu2.get_normalized_args().values()) == 1 assert relu1.get_targets()[0] == list(relu2.get_normalized_args().values())[0] input1 = 1 - node = Node.create_call_cell(Add(), None, ['x'], 'new_conv', - [ScopedValue.create_naming_value('x'), ScopedValue.create_variable_value(input1)], {}, - 'new_conv') + node = Node.create_call_buildin_op(Add(), None, ['x'], 'new_conv', + [ScopedValue.create_naming_value('x'), + ScopedValue.create_variable_value(input1)], {}, + 'new_conv') position = stree.before(relu2) node = stree.insert_node(position, node) # check nodes size