pytorch 模型可视化的例子

yipeiwu_com6年前Python基础

如下所示:

一. visualize.py

from graphviz import Digraph
import torch
from torch.autograd import Variable
 
 
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 torch
from torch.autograd import Variable
from models import *
from visualize import make_dot
x = Variable(torch.rand(1, 3, 256, 256))
model = GeneratorUNet()
y = model(x)
g = make_dot(y)
g.view()

三. 效果展示

以上这篇pytorch 模型可视化的例子就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持【听图阁-专注于Python设计】。

相关文章

Python 分析Nginx访问日志并保存到MySQL数据库实例

使用Python 分析Nginx access 日志,根据Nginx日志格式进行分割并存入MySQL数据库。一、Nginx access日志格式如下:复制代码 代码如下:$remote_...

简单了解Django应用app及分布式路由

简单了解Django应用app及分布式路由

前言 应用在Django的项目中是一个独立的业务模块,可以包含自己的路由,视图,模板,模型. 一 创建应用程序 创建步骤 用manage.py中的子命令startapp创建应用文件夹...

python实现点对点聊天程序

用Python实现点对点的聊天,2个程序,一个是client.py,一个是server.py,通过本机地址127.0.0.1连接进行通信,利用多线程把发送消息和接收消息分开独立进行。 c...

基于Python新建用户并产生随机密码过程解析

说明:本次代码是在Linux下执行的,windows也可以用,把添加用户密码的命令改成windows的就ok了 用Python新建用户并产生随机密码 import passwd_na...

python统计文本字符串里单词出现频率的方法

本文实例讲述了python统计文本字符串里单词出现频率的方法。分享给大家供大家参考。具体实现方法如下: # word frequency in a text # tested wit...