torchrec

Record execution graphs of PyTorch neural networks


License
MIT
Install
pip install torchrec==1.0.0a0

Documentation

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:

Install this package:

$ pip install torchrec

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 during training, you can use TensorBoard and 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.