This package provides a function that generates a dot file from a PyTorch autograd graph.
Saves into result_file
a dot file corresponding to the
autograd graph for the Variable
variable
. The
dictionary variable_labels
associates strings to some
variables, which will be used in the resulting graph.
A typical use is provided in mlp.py:
from torch import nn
from torch.nn import functional as fn
from torch import Tensor
from torch.nn import Module
import agtree2dot
class MLP(Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(MLP, self).__init__()
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
= self.fc1(x)
x = fn.tanh(x)
x = self.fc2(x)
x return x
= MLP(10, 20, 1)
mlp input = Tensor(100, 10).normal_()
= Tensor(100, 1).normal_()
target = mlp(input)
output = nn.MSELoss()
criterion = criterion(output, target)
loss
agtree2dot.save_dot(loss,
{input: 'input',
'target',
target: 'loss',
loss: 'weight1',
mlp.fc1.weight: 'bias1',
mlp.fc1.bias: 'weight2',
mlp.fc2.weight: 'bias2',
mlp.fc2.bias:
},open('./mlp.dot', 'w'))
print('Generated mlp.dot')
try:
'dot', 'mlp.dot', '-Lg', '-T', 'pdf', '-o', 'mlp.pdf' ])
subprocess.check_call([
except subprocess.CalledProcessError:
print('Calling the dot command failed. Is Graphviz installed?')
1)
sys.exit(
print('Generated mlp.pdf')
which would generate a file mlp.dot and try to generate mlp.pdf from it with Graphviz tools.