vilc
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

reinforce subtree in Rewrite 

reinforce Primitive op support in Rewrite
This commit is contained in:
hangangqiang 2022-03-31 14:59:56 +08:00
parent e22b97879c
commit 0e213ca36c
20 changed files with 753 additions and 275 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.jenkins/check/config/filter_pylint.txt
View File

@ -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"

3
mindspore/python/mindspore/rewrite/__init__.py
View File

@ -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"]

4
mindspore/python/mindspore/rewrite/api/node.py
View File

@ -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':
"""

4
mindspore/python/mindspore/rewrite/api/pattern_engine.py
View File

@ -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),

12
mindspore/python/mindspore/rewrite/api/symbol_tree.py
View File

@ -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]:
"""

55
mindspore/python/mindspore/rewrite/api/tree_node_helper.py Normal file
View File

@ -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")

1
mindspore/python/mindspore/rewrite/ast_helpers/ast_finder.py
View File

@ -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.
"""

37
mindspore/python/mindspore/rewrite/ast_helpers/ast_modifier.py
View File

@ -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))

161
mindspore/python/mindspore/rewrite/node.py
View File

@ -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.

130
mindspore/python/mindspore/rewrite/parsers/assign_parser.py
View File

@ -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, "

45
mindspore/python/mindspore/rewrite/parsers/class_def_parser.py
View File

@ -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)

22
mindspore/python/mindspore/rewrite/parsers/module_parser.py
View File

@ -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:

190
mindspore/python/mindspore/rewrite/symbol_tree.py
View File

@ -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:

130
mindspore/python/mindspore/rewrite/symbol_tree_builder.py
View File

@ -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

7
mindspore/python/mindspore/rewrite/topological_manager.py
View File

@ -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=""):
"""

5
tests/ut/python/rewrite/test_net_simple.py
View File

@ -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)

39
tests/ut/python/rewrite/test_node.py
View File

@ -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'),

10
tests/ut/python/rewrite/test_pattern_engine.py
View File

@ -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")

123
tests/ut/python/rewrite/test_subtree_net.py Normal file
View File

@ -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)

49
tests/ut/python/rewrite/test_symbol_tree.py
View File

@ -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