graphviz

conda install graphviz
conda install python-graphviz

import 报错

RuntimeError: failed to execute ['dot', '-Tpdf', '-O', 'test-output/round-table.gv'], make sure the Graphviz executables are on your systems' path

添加环境变量

def print_autograd_graph():
    from graphviz import Digraph
    import torch
    import net
    from torch.autograd import Variable
    import torchvision.models as models

    def make_dot(var, params=None):
        """ Produces Graphviz representation of PyTorch autograd graph

        Blue nodes are the Variables that require grad, orange are Tensors
        saved for backward in torch.autograd.Function

        Args:
            var: output Variable
            params: dict of (name, Variable) to add names to node that
                require grad (TODO: make optional)
        """
        if params is not None:
            #assert all(isinstance(p, Variable) for p in params.values())        
            param_map = {id(v): k for k, v in params.items()}


        node_attr = dict(style='filled',
                        shape='box',
                        align='left',
                        fontsize='12',
                        ranksep='0.1',
                        height='0.2')
        dot = Digraph(node_attr=node_attr, graph_attr=dict(size="12,12"))
        seen = set()

        def size_to_str(size):
            return '('+(', ').join(['%d' % v for v in size])+')'

        def add_nodes(var):
            if var not in seen:
                if torch.is_tensor(var):
                    dot.node(str(id(var)), size_to_str(var.size()), fillcolor='orange')
                elif hasattr(var, 'variable'):
                    u = var.variable
                    #name = param_map[id(u)] if params is not None else ''
                    #node_name = '%s\n %s' % (name, size_to_str(u.size()))
                    node_name = '%s\n %s' % (param_map.get(id(u.data)), size_to_str(u.size()))
                    dot.node(str(id(var)), node_name, fillcolor='lightblue')
                    
                else:
                    dot.node(str(id(var)), str(type(var).__name__))
                seen.add(var)
                if hasattr(var, 'next_functions'):
                    for u in var.next_functions:
                        if u[0] is not None:
                            dot.edge(str(id(u[0])), str(id(var)))
                            add_nodes(u[0])
                if hasattr(var, 'saved_tensors'):
                    for t in var.saved_tensors:
                        dot.edge(str(id(t)), str(id(var)))
                        add_nodes(t)
        add_nodes(var.grad_fn)
        return dot



    torch.manual_seed(1)
    inputs = torch.randn(1,3,224,224)
    model = models.resnet18(pretrained=False)
#    model =net.dehaze_net()
    y = model(Variable(inputs))
    #print(y)

    g = make_dot(y, params=model.state_dict())
    g.view()
    #g

ref
https://zhuanlan.zhihu.com/p/33992733