Visualization of backward pass graph for simple model in PyTorch

visualise-backward-pass.py

#!/usr/bin/evn python3
# Run this script and then the command below.
#
#   dot -Tpng -ograph.png graph.dot
#

import torch as T
from torchviz import make_dot

model = T.nn.Sequential(
    T.nn.Linear(8, 4),
    T.nn.ReLU(),
    T.nn.Linear(4, 1),
)

xs = T.randn((3, 8))
ys = model(xs.requires_grad_())

dot = make_dot(ys, show_attrs=True, show_saved=True)
dot.save('graph.dot')