Computation graph library for Pytorch models, built with Vizstack.
This library allows you to produce a computation graph for Pytorch models that use autograd.
- Import the necessary functions:
from vz_pytorch import start, finish. - Begin recording the graph:
start(model). - Execute your model:
model(inputs). - Finish recording:
graph = finish()
The graph object has __view__() defined to produce the computation graph using Vizstack.