pytoch可视化网络工具-pytorchviz

tensorflow的可视化工具是tensorboard,但是tensorboard在显示网络结构图的时候，显示特别混乱．
之前看了caffe的可视化网络图，很清晰，最近发现pytorch的网络图可视化工具，对网络图的显示也很清晰．

第一步：安装：

一行命令安装
conda install -c anaconda graphviz python-graphviz

第二步：可视化代码（最好自己单独写一个.py文件）
比如我，写了一个net_view.py

from graphviz import Digraph
from torch.autograd import Variable
import torch
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 isinstance(params.values()[0], Variable)
        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()))
                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

第三步：主函数调用上面代码，显示网络

import argparse, os
import torch
import random
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import time
from Net import net_view
def net_graph():
    model = CNN_Net.Gen_Net()
    x = Variable(torch.randn(1, 1, 64, 64))
    x1 = Variable(torch.randn(1, 1, 64, 64))
    x2 = Variable(torch.randn(1, 1, 64, 64))
    x3 = Variable(torch.randn(1, 1, 64, 64))
    y = model(x, x1, x2, x3)#我的网络有４个输入，你自己的网络可以按着自己的情况来
    g = net_view.make_dot(y)
    g.view()

　＃下面是手动输出参数信息![在这里插入图片描述](https://img-blog.csdn.net/20181009155538749?watermark/2/text/aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L0NWX1lPVQ==/font/5a6L5L2T/fontsize/400/fill/I0JBQkFCMA==/dissolve/70)
    params = list(model.parameters())
    k = 0
    for i in params:
        l = 1
        print("该层的结构：" + str(list(i.size())))
        for j in i.size():
            l *= j
        print("该层参数和：" + str(l))
        k = k + l
    print("总参数数量和：" + str(k))

if __name__ == "__main__":
    net_graph()

这样就可以显示了：
显示如下：

参考链接：https://blog.csdn.net/Lucifer_zzq/article/details/80657513
https://github.com/szagoruyko/pytorchviz