torchviz.py文件
import torch
from torch.autograd import Variable
from graphviz import Digraph
def make_dot(var, params=None):
"""
画出 PyTorch 自动梯度图 autograd graph 的 Graphviz 表示.
蓝色节点表示有梯度计算的变量Variables;
橙色节点表示用于 torch.autograd.Function 中的 backward 的张量 Tensors.
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