torchrec 1.0.0a0

Record execution graphs of PyTorch neural networks

torchrec

A small package to record execution graphs of neural networks in PyTorch. The package uses hooks and the grad_fn attribute to record information.
This can be used to generate visualizations at different scope depths.

Licensed under MIT License. View documentation at https://pytorchrec.readthedocs.io/

Installation

Requirements:

• Python3.6+
• PyTorch v1.3 or greater (the cpu version)
• The Graphviz library and graphviz python package.

Install this package:

$ pip install torchrec

Usage

Consider the below example network:

import sys
import torch
import torchrec


class SampleNet(torch.nn.Module):
    def __init__(self):
        torch.nn.Module.__init__(self)

        self.linear_1 = torch.nn.Linear(in_features=3, out_features=3, bias=True)
        self.linear_2 = torch.nn.Linear(in_features=3, out_features=3, bias=True)
        self.linear_3 = torch.nn.Linear(in_features=6, out_features=1, bias=True)
        self.my_special_relu = torch.nn.ReLU()

    def forward(self, inputs):
        x = self.linear_1(inputs)
        y = self.linear_2(inputs)
        z = torch.cat([x, y], dim=1)
        z = self.my_special_relu(self.linear_3(z))
        return z


def main():
    i = int(sys.argv[1])
    net = SampleNet()
    torchrec.render_network(
        net,
        name="Sample Net",
        input_shapes=(1, 3),
        directory="./",
        fmt="svg",
        render_depth=i,
    )


if __name__ == "__main__":
    main()

And visualizations like these can be produced:

Acknowledgements

This is inspired from szagoruyko/pytorchviz. This package differs from pytorchviz as it provides rendering at multiple depths.

Note that for rendering a network via TensorBoard during training, you can use torch.utils.tensorboard.SummaryWriter.add_graph, which records and renders to a protobuf in a single step. The intended usage of pytorchrec is for presentation purposes.