pytorch 模型可视化的例子

野性酷女 2023-01-10 05:04 149阅读 0赞

更多python教程请到： [菜鸟教程][Link 1] https://www.piaodoo.com/

**如下所示：**

一. [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()

三. 效果展示

![Image 1][]

以上这篇pytorch 模型可视化的例子就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持菜鸟教程www.piaodoo.com。

[Link 1]: https://www.piaodoo.com/
[visualize.py]: https://github.com/szagoruyko/functional-zoo/blob/master/visualize.py
[Image 1]: