Exemple #1
0
    def __call__(self, fct, graph=None):
        """Create pydot graph from function.

        Parameters
        ----------
        fct : theano.compile.function.types.Function
            A compiled Theano function, variable, apply or a list of variables.
        graph: pydot.Dot
            `pydot` graph to which nodes are added. Creates new one if
            undefined.

        Returns
        -------
        pydot.Dot
            Pydot graph of `fct`
        """
        if graph is None:
            graph = pd.Dot()

        self.__nodes = {}

        profile = None

        if isinstance(fct, Function):
            profile = getattr(fct, "profile", None)
            fgraph = fct.maker.fgraph
        elif isinstance(fct, gof.FunctionGraph):
            fgraph = fct
        else:
            if isinstance(fct, gof.Variable):
                fct = [fct]
            elif isinstance(fct, gof.Apply):
                fct = fct.outputs
            assert isinstance(fct, (list, tuple))
            assert all(isinstance(v, gof.Variable) for v in fct)
            fgraph = gof.FunctionGraph(inputs=gof.graph.inputs(fct),
                                       outputs=fct)

        outputs = fgraph.outputs
        topo = fgraph.toposort()
        outputs = list(outputs)

        # Loop over apply nodes
        for node in topo:
            nparams = {}
            __node_id = self.__node_id(node)
            nparams["name"] = __node_id
            nparams["label"] = apply_label(node)
            nparams["profile"] = apply_profile(fgraph, node, profile)
            nparams["node_type"] = "apply"
            nparams["apply_op"] = nparams["label"]
            nparams["shape"] = self.shapes["apply"]

            use_color = None
            for opName, color in self.apply_colors.items():
                if opName in node.op.__class__.__name__:
                    use_color = color
            if use_color:
                nparams["style"] = "filled"
                nparams["fillcolor"] = use_color
                nparams["type"] = "colored"

            pd_node = dict_to_pdnode(nparams)
            graph.add_node(pd_node)

            # Loop over input nodes
            for id, var in enumerate(node.inputs):
                var_id = self.__node_id(var.owner if var.owner else var)
                if var.owner is None:
                    vparams = {
                        "name": var_id,
                        "label": var_label(var),
                        "node_type": "input",
                    }
                    if isinstance(var, gof.Constant):
                        vparams["node_type"] = "constant_input"
                    elif isinstance(
                            var, theano.tensor.sharedvar.TensorSharedVariable):
                        vparams["node_type"] = "shared_input"
                    vparams["dtype"] = type_to_str(var.type)
                    vparams["tag"] = var_tag(var)
                    vparams["style"] = "filled"
                    vparams["fillcolor"] = self.node_colors[
                        vparams["node_type"]]
                    vparams["shape"] = self.shapes["input"]
                    pd_var = dict_to_pdnode(vparams)
                    graph.add_node(pd_var)

                edge_params = {}
                if hasattr(node.op, "view_map") and id in reduce(
                        list.__add__, node.op.view_map.values(), []):
                    edge_params["color"] = self.node_colors["output"]
                elif hasattr(node.op, "destroy_map") and id in reduce(
                        list.__add__, node.op.destroy_map.values(), []):
                    edge_params["color"] = "red"

                edge_label = vparams["dtype"]
                if len(node.inputs) > 1:
                    edge_label = str(id) + " " + edge_label
                pdedge = pd.Edge(var_id,
                                 __node_id,
                                 label=edge_label,
                                 **edge_params)
                graph.add_edge(pdedge)

            # Loop over output nodes
            for id, var in enumerate(node.outputs):
                var_id = self.__node_id(var)

                if var in outputs or len(fgraph.clients[var]) == 0:
                    vparams = {
                        "name": var_id,
                        "label": var_label(var),
                        "node_type": "output",
                        "dtype": type_to_str(var.type),
                        "tag": var_tag(var),
                        "style": "filled",
                    }
                    if len(fgraph.clients[var]) == 0:
                        vparams["fillcolor"] = self.node_colors["unused"]
                    else:
                        vparams["fillcolor"] = self.node_colors["output"]
                    vparams["shape"] = self.shapes["output"]
                    pd_var = dict_to_pdnode(vparams)
                    graph.add_node(pd_var)

                    graph.add_edge(
                        pd.Edge(__node_id, var_id, label=vparams["dtype"]))
                elif var.name or not self.compact:
                    graph.add_edge(
                        pd.Edge(__node_id, var_id, label=vparams["dtype"]))

            # Create sub-graph for OpFromGraph nodes
            if isinstance(node.op, builders.OpFromGraph):
                subgraph = pd.Cluster(__node_id)
                gf = PyDotFormatter()
                # Use different node prefix for sub-graphs
                gf.__node_prefix = __node_id
                node.op.prepare_node(node, None, None, "py")
                gf(node.op.fn, subgraph)
                graph.add_subgraph(subgraph)
                pd_node.get_attributes()["subg"] = subgraph.get_name()

                def format_map(m):
                    return str([list(x) for x in m])

                # Inputs mapping
                ext_inputs = [self.__node_id(x) for x in node.inputs]
                int_inputs = [gf.__node_id(x) for x in node.op.local_inputs]
                assert len(ext_inputs) == len(int_inputs)
                h = format_map(zip(ext_inputs, int_inputs))
                pd_node.get_attributes()["subg_map_inputs"] = h

                # Outputs mapping
                ext_outputs = [self.__node_id(x) for x in node.outputs]
                int_outputs = [gf.__node_id(x) for x in node.op.local_outputs]
                assert len(ext_outputs) == len(int_outputs)
                h = format_map(zip(int_outputs, ext_outputs))
                pd_node.get_attributes()["subg_map_outputs"] = h

        return graph
    def __call__(self, fct, graph=None):
        """Create pydot graph from function.

        Parameters
        ----------
        fct : theano.compile.function_module.Function
            A compiled Theano function, variable, apply or a list of variables.
        graph: pydot.Dot
            `pydot` graph to which nodes are added. Creates new one if
            undefined.

        Returns
        -------
        pydot.Dot
            Pydot graph of `fct`
        """
        if graph is None:
            graph = pd.Dot()

        self.__nodes = {}

        profile = None
        if isinstance(fct, Function):
            profile = getattr(fct, "profile", None)
            outputs = fct.maker.fgraph.outputs
            topo = fct.maker.fgraph.toposort()
        elif isinstance(fct, gof.FunctionGraph):
            outputs = fct.outputs
            topo = fct.toposort()
        else:
            if isinstance(fct, gof.Variable):
                fct = [fct]
            elif isinstance(fct, gof.Apply):
                fct = fct.outputs
            assert isinstance(fct, (list, tuple))
            assert all(isinstance(v, gof.Variable) for v in fct)
            fct = gof.FunctionGraph(inputs=gof.graph.inputs(fct), outputs=fct)
            outputs = fct.outputs
            topo = fct.toposort()
        outputs = list(outputs)

        # Loop over apply nodes
        for node in topo:
            nparams = {}
            __node_id = self.__node_id(node)
            nparams['name'] = __node_id
            nparams['label'] = apply_label(node)
            nparams['profile'] = apply_profile(node, profile)
            nparams['node_type'] = 'apply'
            nparams['apply_op'] = nparams['label']
            nparams['shape'] = self.shapes['apply']

            use_color = None
            for opName, color in iteritems(self.apply_colors):
                if opName in node.op.__class__.__name__:
                    use_color = color
            if use_color:
                nparams['style'] = 'filled'
                nparams['fillcolor'] = use_color
                nparams['type'] = 'colored'

            pd_node = dict_to_pdnode(nparams)
            graph.add_node(pd_node)

            # Loop over input nodes
            for id, var in enumerate(node.inputs):
                var_id = self.__node_id(var.owner if var.owner else var)
                if var.owner is None:
                    vparams = {
                        'name': var_id,
                        'label': var_label(var),
                        'node_type': 'input'
                    }
                    if isinstance(var, gof.Constant):
                        vparams['node_type'] = 'constant_input'
                    elif isinstance(
                            var, theano.tensor.sharedvar.TensorSharedVariable):
                        vparams['node_type'] = 'shared_input'
                    vparams['dtype'] = type_to_str(var.type)
                    vparams['tag'] = var_tag(var)
                    vparams['style'] = 'filled'
                    vparams['fillcolor'] = self.node_colors[
                        vparams['node_type']]
                    vparams['shape'] = self.shapes['input']
                    pd_var = dict_to_pdnode(vparams)
                    graph.add_node(pd_var)

                edge_params = {}
                if hasattr(node.op, 'view_map') and \
                        id in reduce(list.__add__,
                                     itervalues(node.op.view_map), []):
                    edge_params['color'] = self.node_colors['output']
                elif hasattr(node.op, 'destroy_map') and \
                        id in reduce(list.__add__,
                                     itervalues(node.op.destroy_map), []):
                    edge_params['color'] = 'red'

                edge_label = vparams['dtype']
                if len(node.inputs) > 1:
                    edge_label = str(id) + ' ' + edge_label
                pdedge = pd.Edge(var_id,
                                 __node_id,
                                 label=edge_label,
                                 **edge_params)
                graph.add_edge(pdedge)

            # Loop over output nodes
            for id, var in enumerate(node.outputs):
                var_id = self.__node_id(var)

                if var in outputs or len(var.clients) == 0:
                    vparams = {
                        'name': var_id,
                        'label': var_label(var),
                        'node_type': 'output',
                        'dtype': type_to_str(var.type),
                        'tag': var_tag(var),
                        'style': 'filled'
                    }
                    if len(var.clients) == 0:
                        vparams['fillcolor'] = self.node_colors['unused']
                    else:
                        vparams['fillcolor'] = self.node_colors['output']
                    vparams['shape'] = self.shapes['output']
                    pd_var = dict_to_pdnode(vparams)
                    graph.add_node(pd_var)

                    graph.add_edge(
                        pd.Edge(__node_id, var_id, label=vparams['dtype']))
                elif var.name or not self.compact:
                    graph.add_edge(
                        pd.Edge(__node_id, var_id, label=vparams['dtype']))

            # Create sub-graph for OpFromGraph nodes
            if isinstance(node.op, builders.OpFromGraph):
                subgraph = pd.Cluster(__node_id)
                gf = PyDotFormatter()
                # Use different node prefix for sub-graphs
                gf.__node_prefix = __node_id
                node.op.prepare_node(node, None, None, 'py')
                gf(node.op.fn, subgraph)
                graph.add_subgraph(subgraph)
                pd_node.get_attributes()['subg'] = subgraph.get_name()

                def format_map(m):
                    return str([list(x) for x in m])

                # Inputs mapping
                ext_inputs = [self.__node_id(x) for x in node.inputs]
                int_inputs = [
                    gf.__node_id(x) for x in node.op.fn.maker.fgraph.inputs
                ]
                assert len(ext_inputs) == len(int_inputs)
                h = format_map(zip(ext_inputs, int_inputs))
                pd_node.get_attributes()['subg_map_inputs'] = h

                # Outputs mapping
                ext_outputs = []
                for n in topo:
                    for i in n.inputs:
                        h = i.owner if i.owner else i
                        if h is node:
                            ext_outputs.append(self.__node_id(n))
                int_outputs = node.op.fn.maker.fgraph.outputs
                int_outputs = [gf.__node_id(x) for x in int_outputs]
                assert len(ext_outputs) == len(int_outputs)
                h = format_map(zip(int_outputs, ext_outputs))
                pd_node.get_attributes()['subg_map_outputs'] = h

        return graph