Introduction

This package provides a function that generates a dot file from a pytorch autograd graph.

Usage

Functions

agtree2dot.save_dot(variable, variable_labels, result_file)

Saves into result_file a dot file corresponding to the autograd graph for variable, which can be either a single Variable or a set of Variables. The dictionary variable_labels associates strings to some variables, which will be used in the resulting graph.

Example

A typical use would be:

import torch

from torch import nn
from torch.nn import functional as fn
from torch import Tensor
from torch.autograd import Variable
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):
        x = self.fc1(x)
        x = fn.tanh(x)
        x = self.fc2(x)
        return x

mlp = MLP(10, 20, 1)
input = Variable(Tensor(100, 10).normal_())
target = Variable(Tensor(100).normal_())
output = mlp(input)
criterion = nn.MSELoss()
loss = criterion(output, target)

agtree2dot.save_dot(loss,
                    { input: 'input', target: 'target', loss: 'loss' },
                    open('./mlp.dot', 'w'))

which would generate a file mlp.dot, which can then be translated to pdf using the Graphviz tools

dot mlp.dot -Lg -T pdf -o mlp.pdf

to produce mlp.pdf.