utils/third_party/nni/common/graph_utils.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.


import logging
import queue
import re
from collections import defaultdict
import torch
from torch.utils.tensorboard._pytorch_graph import NodePy, NodePyIO, NodePyOP, GraphPy
CLASSTYPE_KIND = 'ClassType'
GETATTR_KIND = 'prim::GetAttr'
CAT_KIND = 'aten::cat'
LIST_CONSTRUCT_KIND = 'prim::ListConstruct'
LIST_UNPACK_KIND = 'prim::ListUnpack'
TUPLE_CONSTRUCT_KIND = 'prim::TupleConstruct'
TUPLE_UNPACK_KIND = 'prim::TupleUnpack'

_logger = logging.getLogger(__name__)


def build_module_graph(model, dummy_input):
    return TorchModuleGraph(model, dummy_input)


def build_graph(model, dummy_input, verbose=False):
    g = TorchProtoGraph(model, dummy_input, verbose)
    return g.graph_def, g.stepstats


def parse_traced_name(module_name):
    prefix = 'TracedModule['
    suffix = ']'
    if module_name.startswith(prefix) and module_name.endswith(suffix):
        module_name = module_name[len(prefix):-len(suffix)]
    return module_name


class TorchGraph:
    """
    This class is to extract pytorch model topology graph by tracing
    """

    def __init__(self, model=None, dummy_input=None, traced_model=None):
        """
        Parameters
        ----------
        model : pytorch model
            The model user wants to speed up
        dummy_input : pytorch tensor
            The dummy input for ```jit.trace```, users should put it on right device before pass in
        traced_model : torch._C.torch.jit.TopLevelTracedModule
            An alredy traced model, if traced_model is not None, then TorchGraph will build the graph
            based on this traced model and won't trace the model again.
        """
        assert torch.__version__ >= '1.3.1'
        # check if the input is legal
        if traced_model is not None:
            assert isinstance(traced_model, torch.jit.TopLevelTracedModule)
            self.trace = traced_model
            # it's ok if the graph is already unpacked
            torch._C._jit_pass_inline(self.trace.graph)
        elif model is not None and dummy_input is not None:
            self.bound_model = model
            self._trace(model, dummy_input)
        else:
            raise Exception(
                'Please provide model & dummy_input or the traced_model as inputs')

    def _trace(self, model, dummy_input):
        with torch.onnx.set_training(model, False):
            # import torch.jit
            self.trace = torch.jit.trace(model, dummy_input, check_trace=False)
            torch._C._jit_pass_inline(self.trace.graph)


class TorchProtoGraph(TorchGraph):
    """
    Generates model graph for pytorch models in protobuf, this implementation
    is borrowed from pytorch v1.4.0, and fixed following issues:
    https://github.com/pytorch/pytorch/issues/33691
    https://github.com/pytorch/pytorch/issues/33670

    """

    def __init__(self, model, dummy_input, verbose=False):
        super().__init__(model, dummy_input)

        from tensorboard.compat.proto.config_pb2 import RunMetadata
        from tensorboard.compat.proto.graph_pb2 import GraphDef
        from tensorboard.compat.proto.step_stats_pb2 import StepStats, DeviceStepStats
        from tensorboard.compat.proto.versions_pb2 import VersionDef

        list_of_nodes = self.parse(self.trace.graph, self.trace, dummy_input)
        if verbose:
            print(self.trace.graph)
        self.stepstats = RunMetadata(step_stats=StepStats(
            dev_stats=[DeviceStepStats(device="/device:CPU:0")]))
        self.graph_def = GraphDef(
            node=list_of_nodes, versions=VersionDef(producer=22))

    def parse(self, graph, trace, args=None, omit_useless_nodes=True):
        """This method parses an optimized PyTorch model graph and produces
        a list of nodes and node stats for eventual conversion to TensorBoard
        protobuf format.

        Args:
        graph (PyTorch module): The model graph to be parsed.
        trace (PyTorch JIT TracedModule): The model trace to be parsed.
        args (tuple): input tensor[s] for the model.
        omit_useless_nodes (boolean): Whether to remove nodes from the graph.
        """
        nodes_py = GraphPy()
        for node in graph.inputs():
            if omit_useless_nodes:
                if not node.uses():  # number of user of the node (= number of outputs/ fanout)
                    continue

            if node.type().kind() != CLASSTYPE_KIND:
                nodes_py.append(NodePyIO(node, 'input'))

        attr_to_scope = dict()

        def node_to_name(d):
            return str(d).split(":")[0].strip()
        for node in graph.nodes():
            if node.kind() == GETATTR_KIND:
                attr_name = node.s('name')
                node_name = node_to_name(node)
                parent = node.input().node()
                # If the parent node is not the top-level "self" node
                if parent.kind() == GETATTR_KIND:
                    parent_scope = attr_to_scope[node_to_name(parent)]
                    attr_scope = parent_scope.split('/')[-1]
                    attr_to_scope[node_name] = '{}/{}.{}'.format(
                        parent_scope, attr_scope, attr_name)
                else:
                    attr_to_scope[node_name] = '__module.{}'.format(attr_name)
                # We don't need classtype nodes; scope will provide this information
                if node.output().type().kind() != CLASSTYPE_KIND:
                    node_py = NodePyOP(node)
                    node_py.scopeName = attr_to_scope[node_name]
                    nodes_py.append(node_py)
            else:
                nodes_py.append(NodePyOP(node))

        # Create sink nodes for output ops
        for i, node in enumerate(graph.outputs()):
            node_py = NodePyIO(node, 'output')
            node_py.debugName = "output.{}".format(i + 1)
            node_py.inputs = [node.debugName()]
            nodes_py.append(node_py)

        alias_to_name = dict()
        base_name = parse_traced_name(trace._name)
        for name, module in trace.named_modules(prefix='__module'):
            mod_name = parse_traced_name(module._name)
            attr_name = name.split('.')[-1]
            alias_to_name[name] = '{}[{}]'.format(mod_name, attr_name)

        for node in nodes_py.nodes_op:
            module_aliases = node.scopeName.split('/')[-1].split('.')
            module_name = ''
            for i, alias in enumerate(module_aliases):
                if i == 0:
                    module_name = alias
                    node.scopeName = base_name
                else:
                    module_name += '.' + alias
                    node.scopeName += '/' + \
                        (alias_to_name[module_name]
                         if module_name in alias_to_name else alias)

        nodes_py.populate_namespace_from_OP_to_IO()
        return nodes_py.to_proto()


class NodePyGroup(NodePy):
    """
    This class is used to represent a graph node which consists of multiple jit traced nodes. In a pytorch trace graph,
    there are multiple nodes are traced for one torch.nn.Module object, we group them together to form a single node to
    represent the torch.nn.Module object. We also group some functional call trace nodes together to form a new node.
    """

    def __init__(self, name, unique_name, node_type, op_type, node_cpps, inputs=None, outputs=None, key_node=None):
        """
        Parameters:
        -----------
        name: str
            node name, such as `conv1`, `backbone.classifier`
        unique_name: str
            A global unique name for current node. Due to some modules,
            such as relu, may be reused several times, so the scopename
            is not suitable as the global unique identifier, so we add a
            unique_name for each node as the global unique identifier.
            We should use the unique_name to traverset the module graph.
        node_type: str
            `module` or `func`
        op_type: str
            operation type, such as `Conv2d`, `aten::view`
        node_cpps: list of torch._C.Node
            jit trace nodes which are included in this new node
        inputs: list of str
            All the inputs of this node, each element is debugName of one input
        outputs: list of str
            All the outputs of this node, each element is debugName of one output
        key_node: torch._C.Node
            The key node of this NodePyGroup.
        """
        super(NodePyGroup, self).__init__(name, [])
        self.node_cpps = node_cpps
        self.name = name
        self.unique_name = unique_name
        self.op_type = op_type
        self.type = node_type
        self.nodes = []
        self.auxiliary = None
        self.add_nodes(node_cpps)
        self.inputs = inputs
        self.outputs = outputs
        # The core node in this NodePyGroup
        self.key_node = key_node

    def add_nodes(self, node_cpps):
        for node_cpp in node_cpps:
            nodepy = NodePyOP(node_cpp)
            nodepy.name = node_cpp.scopeName() + '_' + node_cpp.kind()
            self.nodes.append(nodepy)

    def sub_node_names(self):
        return [x.name for x in self.nodes]

    def __repr__(self):
        return 'name: {}, type: {}, op_type: {}, sub_nodes: {}, inputs: {}, outputs: {}, aux: {}'.format(
            self.name, self.type, self.op_type, self.sub_node_names(),
            self.inputs, self.outputs, self.auxiliary
        )


class TorchModuleGraph(TorchGraph):
    """
    Generates model graph, each node is created from single or multiple jit trace nodes.
    """

    def __init__(self, model=None, dummy_input=None, traced_model=None):
        super().__init__(model, dummy_input, traced_model)
        self.global_count = 0
        self.name_to_node, self.input_to_node, self.output_to_node = self._build_graph()
        self._extract_auxiliary_info()

    def _expand_key_func_node(self, node, nodes, input_to_node, output_to_node,
                              module_type):
        """
        For trace graph nodes, some nodes are not in modules, these nodes are usually generated by
        the functions directly called in module ```forward```. For such nodes, some of them are
        trivial op which are label by ```prim::```, some of them are not such ops which is call
        non-prim ops. This function is to merge neighbor prim ops to a non-prim op, to construct
        a node.

        Parameters
        ----------
        node : trace graph node
            The non-prim node to expand
        nodes : list of trace graph node
            All the trace graph nodes within the same scope as the non-prim node
        input_to_node : dict
            key: input name, value: a node that uses this input
        output_to_node : dict
            key: output name, value: a node that generates this output
        module_type : str
            can be 'module' or 'func'

        Returns
        -------
        node
            the expanded non-prim node
        """
        # TODO: scope name could be empty
        node_name = '.'.join([self._get_module_name(
            node.scopeName()), node.kind(), str(self.global_count)])
        unique_name = node_name
        _logger.debug("expand non-prim node, node name: %s", node_name)
        self.global_count += 1
        op_type = node.kind()
        node_group = [node]
        inputs = list()
        outputs = list()
        node_queue = queue.Queue()
        node_queue.put(node)
        while not node_queue.empty():
            curr_node = node_queue.get()
            for _input in curr_node.inputs():
                input_name = _input.debugName()
                if input_name in output_to_node and output_to_node[input_name] in nodes:
                    predecessor_node = output_to_node[input_name]
                    if not self._is_key_func(predecessor_node):
                        node_group.append(predecessor_node)
                        node_queue.put(predecessor_node)
                    else:
                        inputs.append(input_name)
                else:
                    inputs.append(input_name)
        for output in node.outputs():
            outputs.append(output.debugName())
        nodepy = NodePyGroup(node_name, unique_name, module_type, op_type,
                             node_group, inputs=inputs, outputs=outputs, key_node=node)
        return nodepy

    def _expand_module_node(self, node, node_name, unique_name, op_type, nodes,
                            input_to_node, output_to_node, module_type):
        """
        merge the adjacent nodes of the module. The difference between the
        _expand_module_node and _expand_non_prim_node is that, the _expand_non_prim_node
        only merge the prim:: nodes into the aten:: node, in contrast,the _expand_module_node
        will merge all adjacent nodes into a same nodepy group.

        Parameters
        ----------
        node : trace graph node
            The non-prim node to expand
        node_name : str
            specify the node_name for NodePyGroup
        unique_name : str
            unique_name for the NodePyGroup
        op_type : str
            specify the op_type for the NodePyGroup
        nodes : list of trace graph node
            All the trace graph nodes within the same scope as the non-prim node
        input_to_node : dict
            key: input name, value: a node that uses this input
        output_to_node : dict
            key: output name, value: a node that generates this output
        module_type : str
            can be 'module' or 'func'
        Returns
        -------
        node
            the expanded non-prim node

        """
        _logger.debug("expand module node, node name: %s", node_name)
        self.global_count += 1
        if not op_type:
            op_type = node.kind()
        node_group = [node]
        inputs = list()
        outputs = list()
        node_queue = queue.Queue()
        node_queue.put(node)
        visited = {node}
        while not node_queue.empty():
            curr_node = node_queue.get()
            for _input in curr_node.inputs():
                input_name = _input.debugName()
                if input_name in output_to_node and output_to_node[input_name] in nodes:
                    predecessor_node = output_to_node[input_name]
                    if predecessor_node not in visited:
                        node_group.append(predecessor_node)
                        node_queue.put(predecessor_node)
                        visited.add(predecessor_node)
                else:
                    inputs.append(input_name)
            for _output in curr_node.outputs():
                output_name = _output.debugName()
                if output_name in input_to_node and input_to_node[output_name] in nodes:
                    successor_node = input_to_node[output_name]
                    if successor_node not in visited:
                        node_group.append(successor_node)
                        node_queue.put(successor_node)
                        visited.add(successor_node)
                else:
                    outputs.append(output_name)

        nodepy = NodePyGroup(node_name, unique_name, module_type, op_type,
                             node_group, inputs=inputs, outputs=outputs)
        return nodepy

    def _extract_cat_info(self, node_group, cpp_node):
        """
        Extract the detail information of the cat operation,
        such the order of the input tensor, the shape of each
        input tensor, the output shape, and the cat dimension.

        Parameters
        ----------
        node_group : NodePyGroup
        cpp_node: torch._C.Node
            It should be ```aten::cat``` node

        Returns
        -------
        dict
            Include auxiliary information for the cat operation.
            This dict objec has four keys: 'cat_dim', 'out_shape',
            'in_order' and 'in_shape'. cat_dim is the dimension of
            the cat operation to concat the input tensors. out_shape
            is the shape of the output tensor of the cat operation.
            in_order is an ordered list which contains the corresponding
            parent operaion nodes of the input tensors. in_shape is also
            an ordered list that contains the input shapes of the input
            tensor.
        """
        # only suport the cat operation
        assert cpp_node.kind() == CAT_KIND
        cat_info = {}
        # get the shape of the output tensor
        t_output = cpp_node.output()
        out_shape = t_output.type().sizes()
        cat_info['out_shape'] = out_shape
        # get the cat dimension
        inputs = cpp_node.inputs()
        cat_dim = list(inputs)[1].toIValue()
        cat_info['cat_dim'] = cat_dim
        # get the order of the input tensors
        # To get the order of the input tensors, we need
        # to be aware of the topology of the model, which
        # means we should extract the auxiliary information
        # after the build_index function.
        input_order = []
        list_construct_cpp = list(cpp_node.inputs())[0].node()
        input_tensors = list(list_construct_cpp.inputs())
        for _tensor in input_tensors:
            debug_name = _tensor.debugName()
            input_order.append(self.output_to_node[debug_name].unique_name)
        cat_info['in_order'] = input_order
        input_shapes = [t.type().sizes() for t in input_tensors]
        cat_info['in_shape'] = input_shapes
        return cat_info

    def _extract_linear_shape_info(self, node_group):
        """
        Extract linear shape input/output tensor shape info from its aten::addmm op.

        Parameters
        ----------
        node_group : NodePyGroup
            NodePyGroup object associated with the linear module.

        Returns
        -------
        dict
            Include shape of input tensor and shape of output tensor
        """
        for cpp_node in node_group.node_cpps:
            if cpp_node.kind() == 'aten::addmm':
                # https://github.com/pytorch/pytorch/blob/1.6/torch/nn/functional.py#L1682
                # inputs of aten::addmm:
                # inputs[0] is bias
                # inputs[1] is input data
                # inputs[2] is weight
                t_input = list(cpp_node.inputs())[1]
                t_output = cpp_node.output()
                assert isinstance(t_input.type(), torch._C.TensorType)
                assert isinstance(t_output.type(), torch._C.TensorType)
                in_shape = t_input.type().sizes()
                out_shape = t_output.type().sizes()
                return {'in_shape': in_shape, 'out_shape': out_shape}
        return None

    def _extract_shape_info(self, node):
        """
        Extract the shape information of ```aten::view``` node

        Parameters
        ----------
        node : trace graph node
            It should be ```aten::view``` node

        Returns
        -------
        dict
            Include shape of input tensor and shape of output tensor
        """
        t_input = None
        for _input in node.inputs():
            t_input = _input
            break
        t_output = node.output()
        assert isinstance(t_input.type(), torch._C.TensorType)
        assert isinstance(t_output.type(), torch._C.TensorType)
        in_shape = t_input.type().sizes()
        out_shape = t_output.type().sizes()
        return {'in_shape': in_shape, 'out_shape': out_shape}

    def _extract_leaf_modules(self):
        """
        Extract leaf modules from the given graph. Leaf module means it does not have submodules.
        To extract leaf modules because only leaf module can be replaced. And shape inference can
        be done in leaf module level. Other shape inference is done in lower level i.e.,
        operation level.

        Returns
        -------
        list
            a list of scope name of all the leaf modules
        """
        def is_parent(name1, name2):
            """
            check if name1 is parent node of name2, for example:
            name1: aa.bb,  name2: aa.bb.cc,  return True
            name1: aa.b,  name2: aa.bb, return False
            """
            parts1, parts2 = name1.split('.'), name2.split('.')
            if len(parts1) >= len(parts2):
                return False
            for i, _ in enumerate(parts1):
                if parts2[i] != parts1[i]:
                    return False
            return True
        module_names = sorted([x[0]
                               for x in self.trace.named_modules() if x[0]])
        leaf_nodes = []
        for i, name in enumerate(module_names):
            if i + 1 >= len(module_names) or not is_parent(name, module_names[i + 1]):
                leaf_nodes.append(name)
        return leaf_nodes

    def _get_module_name(self, scope_name):
        """
        Retrieve module name from scope name.
        Parameters:
        -----------
        scope_name: str
            scope_name of a graph node, for example:
            for pytorch 1.3.1: MyModel/BackboneModel[backbone]/Conv2d[conv2]
            for pytorch 1.4.0: __module.backbone/__module.backbone.conv2

        Returns:
        -------
        str
            module name, such as backbone.conv2
        """
        if torch.__version__ >= '1.4.0':
            return scope_name.split('/')[-1].replace('__module.', '')
        else:
            return '.'.join(re.findall(r'\[(.*?)\]', scope_name))

    def _build_index(self, nodes_op):
        name_to_node = dict()
        input_to_node = defaultdict(list)
        output_to_node = dict()
        for node in nodes_op:
            name_to_node[node.unique_name] = node
            for _input in node.inputs:
                input_to_node[_input].append(node)
            for output in node.outputs:
                assert not output in output_to_node, \
                    "One output cannot be generated by multiple nodes"
                output_to_node[output] = node
        return name_to_node, input_to_node, output_to_node

    def _is_key_func(self, node_cpp):
        """
        Judge if a cpp node is a key function node.
        If so, we should not merge this node into the
        adjacent node.
        """
        if node_cpp.kind().startswith('aten::'):
            # the nodes that start with 'aten' are key function
            # nodes
            return True
        if node_cpp.kind() in [LIST_UNPACK_KIND, TUPLE_UNPACK_KIND]:
            # We cannot merge the List/Tuple
            # Unpack func into other nodes, else it
            # may lead to a graph construction error.
            # The reason why we donnot take the construct node
            # also as a key node is that `cat` operation node need
            # the last(previous) visited node to infer the mask. If
            # we take the Construct node as the important node, the
            # predecessor of the `cat` node will always be a construct
            # node, which means we cannot infer the mask for the cat
            # operation.
            return True
        return False

    def unpack_manually(self):
        """
        Unpack the tensor tuple or tensor list manually,
        and remove the ListUnpack/TupleUnpack node from
        the graph. Note: this function will change the
        graph structure.
        """
        if hasattr(self, 'unpacked'):
            # if already unpacked the tuple/list manually
            return
        for node in self.nodes_py.nodes_op:
            if node.op_type in [TUPLE_UNPACK_KIND, LIST_UNPACK_KIND]:
                unpack_cpp = node.key_node
                last_cpp = list(unpack_cpp.inputs())[0].node()
                if last_cpp.kind() in [TUPLE_CONSTRUCT_KIND, LIST_CONSTRUCT_KIND]:
                    # we need check if the tensor tuple or tensor list is produced
                    # by a list/tuple construct node. If so, we can unpack the tuple
                    # or list manunally.
                    _logger.debug('List/Tuple Construct Node(cpp) %s', str(last_cpp))
                    _logger.debug('List/Tuple Unpack Node(cpp) %s', str(unpack_cpp))
                    assert len(list(unpack_cpp.outputs())) == len(list(last_cpp.inputs()))
                    errmsg = '%s Input number: %d if inconsistent with the output number %d' % (unpack_cpp, \
                        len(node.inputs), len(list(last_cpp.inputs())))

                    assert len(node.inputs) == len(list(last_cpp.inputs())), errmsg
                    for _debug_input, _debug_output in zip(node.inputs, node.outputs):
                        # _debug_input = _input.debugName()
                        # _debug_output = _output.debugName()
                        if _debug_input in self.input_to_node and _debug_output in self.input_to_node:
                            # input_to_node[_debug_input] is a list of NodePyGroup, because
                            # one tensor can be used as input for multiple nodes at the same time.

                            # note that, in this case, the construct cpp node and unpack cpp node
                            # will be merged into the same NodePyGroup, so we remove the `node` from
                            # input_to_node[_debug_input] and directly connect this tensor to the
                            # input_to_node[_debug_output]
                            self.input_to_node[_debug_input].remove(node)
                            # add the following nodes of _output into the input_to_node[_debug_input]
                            self.input_to_node[_debug_input].extend(self.input_to_node[_debug_output])
                        # just remove the _debug_output from the grapgh index. So that we can also skip
                        # the construct and tuple
                        if _debug_output in self.input_to_node:
                            for following_node in self.input_to_node[_debug_output]:
                                _tmp_index = following_node.inputs.index(_debug_output)
                                following_node.inputs[_tmp_index] = _debug_input


        self.unpacked = True

    def _build_graph(self):
        """
        Build graph using our defined format from jit trace.
        There are basically three steps: first, construct necessary information (data structures),
        second, extract all the modules to convert to node, Third, extract all functions to convert
        to node.

        Returns
        -------
        dict
            use name to index nodes, key: node name, value: node
        dict
            use input (its name) to index nodes,
            key: input, value: list of nodes that take this input
        dict
            use output (its name) to index nodes,
            key: output, value: node that generates this output
        """
        omit_useless_nodes = True
        graph = self.trace.graph
        # _logger.debug(graph)
        # build output mapping, from output debugName to its node
        output_to_node = {x.debugName(): n for n in graph.nodes()
                          for x in n.outputs()}
        # build input mapping, from input debugName to its node
        input_to_node = {x.debugName(): n for n in graph.nodes()
                         for x in n.inputs()}
        # build module mapping, from module name to all nodes (as list) under this module scope
        module_to_nodes = defaultdict(list)
        # the mapping of function (non-module in forward) to nodes, key is scope name
        func_to_nodes = defaultdict(list)

        nodes_py = GraphPy()
        for node in graph.inputs():
            if omit_useless_nodes:
                if not node.uses():  # number of user of the node (= number of outputs/ fanout)
                    continue

            if node.type().kind() != 'ClassType':
                nodes_py.append(NodePyIO(node, 'input'))

        self.leaf_modules = self._extract_leaf_modules()
        module_to_type = {name: parse_traced_name(
            module._name) for name, module in self.trace.named_modules()}

        # associate module name with their trace graph nodes
        for node in graph.nodes():
            module_name = self._get_module_name(node.scopeName())
            if module_name in self.leaf_modules:
                module_to_nodes[module_name].append(node)
            else:
                func_to_nodes[node.scopeName()].append(node)
        # build node group for module
        for module_name, node_cpps in module_to_nodes.items():
            use_count = 0
            merged = set()
            for node in node_cpps:
                if node not in merged:
                    # modules that have same scope name may have different locations in the
                    # graph. Futhermore, there are also lots of prim:: nodes that in node_cpps,
                    # so we also need to call the expand_module_node.
                    unique_name = module_name
                    if use_count > 0:
                        unique_name = module_name + '.%d' % use_count
                    node_group = self._expand_module_node(
                        node, module_name, unique_name, module_to_type[module_name],
                        node_cpps, input_to_node, output_to_node, 'module')
                    nodes_py.nodes_op.append(node_group)
                    use_count += 1
                    merged.update(node_group.node_cpps)

        # each scope_name may have multiple funcs, we split them and create node for each of them
        # build node group for torch.nn.functional
        for _, nodes in func_to_nodes.items():
            # extract non prim:: nodes
            key_func_nodes = list()
            for node in nodes:
                if self._is_key_func(node):
                    # find the key function nodes
                    key_func_nodes.append(node)
            # for each non prim node, expand it
            for node in key_func_nodes:
                node_group = self._expand_key_func_node(
                    node, nodes, input_to_node, output_to_node, 'func')
                nodes_py.nodes_op.append(node_group)
                # get shape infor for view (aten::view) func
                # if node_group.op_type in ['aten::view', 'aten::flatten']:
                #     node_group.auxiliary = self._extract_shape_info(node)

        for node in graph.outputs():  # Create sink nodes for output ops
            node_py = NodePyIO(node, 'output')
            nodes_py.append(node_py)

        self.nodes_py = nodes_py
        # build index
        return self._build_index(self.nodes_py.nodes_op)

    def _extract_auxiliary_info(self):
        """
        Extract the auxiliary information for the nodegroups
        if necessary. For example, view/flatten operations may
        need the shape of the input tensor and output tensor.
        """
        # extract the input & output shape for the view and flatten
        for node_group in self.nodes_py.nodes_op:
            if node_group.op_type in ['aten::view', 'aten::flatten', 'aten::mean', 'aten::reshape']:
                # get shape infor for view (aten::view) func
                cpp_node = list(filter(lambda x: x.kind() == node_group.op_type,
                                       node_group.node_cpps))[0]
                node_group.auxiliary = self._extract_shape_info(cpp_node)
            elif node_group.op_type == 'Linear':
                node_group.auxiliary = self._extract_linear_shape_info(node_group)
            elif node_group.op_type == CAT_KIND:
                # get the detail information for cat func
                cpp_node = list(filter(lambda x: x.kind() == node_group.op_type,
                                       node_group.node_cpps))[0]
                node_group.auxiliary = self._extract_cat_info(
                    node_group, cpp_node)

    def find_predecessors(self, unique_name):
        """
        Find predecessor node of the given node

        Parameters
        ----------
        unique_name : str
            The unique name of the node

        Returns
        -------
        list
            a list of nodes who are the given node's predecessor
        """
        predecessors = []
        for _input in self.name_to_node[unique_name].inputs:
            if not _input in self.output_to_node:
                _logger.debug("cannot find node with %s as its output", _input)
            else:
                node_py = self.output_to_node[_input]
                predecessors.append(node_py.unique_name)
        return predecessors

    def find_successors(self, unique_name):
        """
        Find successor nodes of the given node

        Parameters
        ----------
        unique_name : str
            The unique name of the node

        Returns
        -------
        list
            a list of nodes who are the given node's successor
        """
        successors = []
        for output in self.name_to_node[unique_name].outputs:
            if output not in self.input_to_node:
                # may reach the output of the whole graph
                continue
            nodes_py = self.input_to_node[output]
            for node_py in nodes_py:
                successors.append(node_py.unique_name)
        return successors